diff --git a/CHANGELOG.md b/CHANGELOG.md index 18e53bc..c6624db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,13 +9,54 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased -- Normalised `hash_element_at` so NaN maps to a dummy/sentinel value +### Added + +**NdArray, XArray and DLPack Feature**: +- `NdArray` n-dimensional dense array (`ndarray` feature) over `Buffer`, + generic over f32/f64, with NaN null semantics, `NdArrayV` zero-copy views, + transpose and axis permutation, and Table/Matrix/Array interop. +- `SuperNdArray` chunked n-dimensional array with `SuperNdArrayV` + chunk-spanning views and rechunking. +- `XArray` labelled n-dimensional array (`xarray` feature) with named + dimensions, coordinate selection via `at`/`between`/`nearest`, and owned, + view, or chunked storage behind one type. +- DLPack tensor interchange (`dlpack` feature) over the legacy and 1.x + versioned ABIs for zero-copy sharing with PyTorch, JAX, and TensorFlow. +- `AxisSelection` trait for `.s()` axis slicing, with `NdArrayVT`, + `SuperNdArrayVT`, and `XArrayVT` view tuple aliases. +- Broadcast kernels for `NdArray`, `SuperNdArray`, and `XArray`, plus + `Value` variants, conversions, and concatenation for the new types. +- `NdArray` in minarrow-py with the DLPack protocol - `__dlpack__`, + `from_dlpack`, and named bridges to NumPy, PyTorch, JAX, TensorFlow, + and CuPy. +- `PyNdArray` in minarrow-pyo3 (`ndarray` feature) for Rust extensions + bridging tensors to Python, with the DLPack capsule glue hosted in + its `ffi::dlpack` and shared by minarrow-py. + +**Minor Feature Additions** - At `value_at` to `Array` and `ArrayV` for (opt-in) normalised equality checking. -- Normalised equality checking for `Scalar` -0.0 == 0.0 and NaN == NaN. -- Fixed null handling bug on String and Categorical set_str method. +- Added bitmask-driven gather to Array, Table and their view variants. +- Added default SIMD chipset behaviour. +- Added bitwise kernels. +- Added default display trait impl to `Scalar`. +- Added `LBuffer::freeze()`. +- Added `NdArray` n-dimensional array with views and chunked variants. +- Added `XArray` labelled n-dimensional array. +- Added axis selection, broadcasting, and `Value` support for the n-dimensional types. +- Added DLPack FFI for zero-copy PyTorch, JAX, and TensorFlow interchange, with `NdArray` in minarrow-py and minarrow-pyo3. - Added `get`, `get_unchecked`, `get_str` and `get_str_unchecked` methods to `Array` - Added bitmask gather to Array, Table and their view variants. +### Bugfixes +- Fixed a kernel branching issue in the arithmetic SIMD paths. +- Fixed out-of-range categorical codes to resolve to the empty string. +- Fixed `StringArray` length reporting the byte length when `MaskedArray` was not imported. +- Fixed null handling bug on String and Categorical set_str method. + +### Modified +- Bumped pyo3 to 0.29. +- Normalised `hash_element_at` so NaN maps to a dummy/sentinel value +- Normalised equality checking for `Scalar` -0.0 == 0.0 and NaN == NaN. ## [0.15.0] - 2026-06-30 ### Added diff --git a/Cargo.toml b/Cargo.toml index 49f6fa9..b218271 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -138,6 +138,24 @@ views = [] # with large data sizes. matrix = [] +# N-dimensional f32/f64 container with compact column-major layout. +# Slicing and axis selection are available through the `select` feature. +# Useful for bridging into Python ML/AI runtimes via DLPack. +# +# Note that this does not include statistical routines. +ndarray = [] + +# DLPack tensor interchange with Python and other runtimes. Compatible +# ownership, alignment, layout, and protocol choices allow direct sharing. +# Requires the `ndarray` feature. +dlpack = ["ndarray"] + +# Labelled N-dimensional array wrapping NdArray with named dimensions +# and optional coordinate labels per axis. Requires `ndarray` feature. +# Coordinate value lookup is available when `scalar_type`, `views`, and +# `select` are also enabled. +xarray = ["ndarray"] + # Adds `to_apache_arrow()` for casting into that library. cast_arrow = ["arrow", "arrow-schema"] @@ -238,7 +256,7 @@ fast_dict = [ "dep:parking_lot", ] -# Adds typed arithmetic broadcasting for add, sub, mult, div, rem +# Adds typed arithmetic broadcasting for add, sub, mult, div, rem. broadcast = [] # Adds Hash and Eq implementations for Scalar, and hash_element_at for Array. @@ -312,33 +330,49 @@ rustdoc-args = ["--cfg", "docsrs"] [[example]] name = "arithmetic" path = "examples/arithmetic.rs" -required-features = ["broadcast"] +required-features = ["broadcast", "value_type"] # Broadcasting examples [[example]] name = "test_broadcasting" path = "examples/broadcasting/test_broadcasting.rs" -required-features = ["broadcast"] +required-features = ["broadcast", "value_type"] [[example]] name = "test_scalar_arithmetic" path = "examples/broadcasting/test_scalar_arithmetic.rs" -required-features = ["broadcast", "scalar_type"] +required-features = ["broadcast", "value_type", "scalar_type"] [[example]] name = "test_string_broadcasting" path = "examples/broadcasting/test_string_broadcasting.rs" -required-features = ["broadcast"] +required-features = ["broadcast", "value_type"] [[example]] name = "test_value_ops" path = "examples/broadcasting/test_value_ops.rs" -required-features = ["broadcast"] +required-features = ["broadcast", "value_type"] [[example]] name = "test_value_macros" path = "examples/broadcasting/test_value_macros.rs" -required-features = ["broadcast"] +required-features = ["broadcast", "value_type"] + +# N-dimensional examples +[[example]] +name = "ndarray" +path = "examples/ndarray.rs" +required-features = ["ndarray"] + +[[example]] +name = "super_ndarray" +path = "examples/super_ndarray.rs" +required-features = ["ndarray"] + +[[example]] +name = "xarray" +path = "examples/xarray.rs" +required-features = ["xarray", "scalar_type", "views", "select"] [[bench]] name = "hotloop_benchmark_simd" diff --git a/README.md b/README.md index e4c6a2b..8a60a0a 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,8 @@ Semantic groupings (`NumericArray`, `TextArray`, `TemporalArray`) support flexib Bonus LAPACK-compatible `Matrix` and `Cube` types support analytical workload variations. +`NdArray` n-dimensional numeric data, with chunked (`SuperNdArray`) and labelled (`XArray`) forms, and hands tensors to PyTorch, JAX, and NumPy zero-copy over DLPack. + ### Zero-Copy Views Ergonomic zero-copy row and column selection. @@ -225,7 +227,7 @@ Minarrow's direct access is within the noise threshold of raw Vec performance wh - **Rust to Python¹** 1m rows, 2 columns : **165ns** - **Python to Rust¹** 1m rows, 2 columns: **2.5μs** -See `minarrow-pyo3/examples` +See `minarrow-pyo3/examples` ### Test machine ▎ ¹ = Intel Core Ultra 7 155H · 32 GB · Ubuntu 24.04 · 1.97-nightly release build. @@ -262,6 +264,9 @@ Additional types: | `value_type` | Catch-all `Value` enum for unified typing | | `matrix` | 2D matrix with BLAS/LAPACK-compatible layout | | `cube` | Stacks tables along an extra axis | +| `ndarray` | N-dimensional array for bridging into ML & AI, with a chunked version and native DLPack interchange | +| `xarray` | N-dimensional array with named dims and coordinates for spatial queries | +| `dlpack` | DLPack tensor interchange with PyTorch, JAX, TensorFlow | | `shared_dict` | Shared source of truth for categorical dictionaries. | Performance: diff --git a/examples/ndarray.rs b/examples/ndarray.rs new file mode 100644 index 0000000..da07f9c --- /dev/null +++ b/examples/ndarray.rs @@ -0,0 +1,71 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use minarrow::traits::selection::RowSelection; +use minarrow::{Concatenate, NdArray, nd}; + +fn main() { + // Flat input follows NdArray's column-major layout. + let a = NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0], + &[4, 3], + ); + + println!("=== Shape and access ===\n"); + println!("shape: {:?}", a.shape()); + println!("strides: {:?}", a.strides()); + println!("n_obs: {}", a.n_obs()); + println!("a.get(&[2, 1]) = {}\n", a.get(&[2, 1])); + + // Constructors create ordinary NdArray containers. + println!("=== Constructors ===\n"); + let steps = NdArray::::linspace(0.0, 1.0, 5); + println!("linspace(0.0, 1.0, 5): {:?}", steps.as_slice()); + let identity = NdArray::::eye(3); + println!("eye(3) diagonal: {:?}\n", (0..3).map(|i| identity.get(&[i, i])).collect::>()); + + // Views can change shape, offsets, or strides without copying the data. + println!("=== Views ===\n"); + println!("a.obs(1) - one observation as a 1D view"); + let row = a.obs(1); + println!(" shape {:?}, values {:?}\n", row.shape(), (0..3).map(|j| row.get(&[j])).collect::>()); + + println!("a.slice(nd![1..3, 0..2]) - range on both axes"); + let window = a.slice(nd![1..3, 0..2]); + println!(" shape {:?}, window[0, 0] = {}\n", window.shape(), window.get(&[0, 0])); + + println!("a.r(0..2) - axis-0 selection"); + let head = a.r(0..2); + println!(" shape {:?}\n", head.shape()); + + println!("view.transpose() - zero-copy stride swap"); + let t = a.as_view().transpose(); + println!(" shape {:?}, t.get(&[1, 2]) = {}\n", t.shape(), t.get(&[1, 2])); + + // Apply operation to each array element. + println!("=== Apply ===\n"); + let doubled = a.apply(|v| v * 2.0); + println!("a.apply(|v| v * 2.0): doubled[3, 2] = {}\n", doubled.get(&[3, 2])); + + // Concatenation extends axis 0 while retaining the trailing shape. + println!("=== Concatenate ===\n"); + let b = NdArray::from_slice(&[13.0, 14.0, 15.0], &[1, 3]); + let joined = a.clone().concat(b).unwrap(); + println!("a.concat(b): shape {:?}, last row starts with {}\n", joined.shape(), joined.get(&[4, 0])); + + // A 2D conversion produces one Table column per axis-1 entry. + println!("=== To Table ===\n"); + let table = a.clone().to_table(None).unwrap(); + println!("{}", table); +} diff --git a/examples/super_ndarray.rs b/examples/super_ndarray.rs new file mode 100644 index 0000000..6d98d99 --- /dev/null +++ b/examples/super_ndarray.rs @@ -0,0 +1,59 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use minarrow::structs::chunked::super_array::RechunkStrategy; +use minarrow::{Consolidate, NdArray, SuperNdArray}; + +fn main() { + // Simulate three sensor batches arriving with different observation + // counts. Every batch has two measurements per observation. + let mut snd = SuperNdArray::new("sensor_frames"); + snd.push(NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2])); + snd.push(NdArray::from_slice(&[3.0, 4.0, 5.0, 30.0, 40.0, 50.0], &[3, 2])); + snd.push(NdArray::from_slice(&[6.0, 60.0], &[1, 2])); + + println!("=== Shape across batches ===\n"); + println!("n_batches: {}", snd.n_batches()); + println!("n_obs: {}", snd.n_obs()); + println!("shape: {:?}\n", snd.shape()); + + // Global indices address the combined observation range across batches. + println!("=== Global access ===\n"); + println!("snd.get(&[0, 0]) = {} (batch 0)", snd.get(&[0, 0])); + println!("snd.get(&[3, 1]) = {} (batch 1)", snd.get(&[3, 1])); + println!("snd.get(&[5, 1]) = {} (batch 2)\n", snd.get(&[5, 1])); + + // The requested observation window starts in one batch and ends in the next. + println!("=== Batch-spanning window ===\n"); + let window = snd.slice(1, 3); + println!("snd.slice(1, 3): n_obs {}, spans {} slices", window.n_obs(), window.n_slices()); + println!("window.get(&[1, 0]) = {}\n", window.get(&[1, 0])); + + // Consolidate only when a consumer requires one contiguous allocation. + println!("=== Consolidate ===\n"); + let flat = snd.clone().consolidate(); + println!("consolidated shape: {:?}", flat.shape()); + println!("flat.get(&[5, 1]) = {}\n", flat.get(&[5, 1])); + + // Change the batch boundaries without changing the logical values. + println!("=== Rechunk ===\n"); + let mut even = snd.clone(); + even.rechunk(RechunkStrategy::Count(2)).unwrap(); + println!("rechunk(Count(2)): n_batches {}, first batch obs {}", even.n_batches(), even.batch(0).unwrap().shape()[0]); + + // Batch boundaries do not affect logical equality. + println!("\n=== Logical equality ===\n"); + let single = SuperNdArray::from_batches(vec![flat], "sensor_frames"); + println!("snd == consolidated single batch: {}", snd == single); +} diff --git a/examples/xarray.rs b/examples/xarray.rs new file mode 100644 index 0000000..acf3e2d --- /dev/null +++ b/examples/xarray.rs @@ -0,0 +1,66 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use minarrow::{NdArray, XArray, arr_f64, arr_str32}; + +fn main() { + // Four hourly observations of three measurements. + let data = NdArray::from_slice( + &[ + 20.1, 20.4, 20.9, 21.3, // temp + 1.01, 1.02, 1.00, 0.99, // pressure + 55.0, 54.0, 52.0, 51.0, // humidity + ], + &[4, 3], + ); + let mut xa = XArray::new(data, &["hour", "measurement"]); + xa.assign_coords("hour", arr_f64![0.0, 1.0, 2.0, 3.0]); + + println!("=== Dims and coords ===\n"); + println!("{:?}", xa); + println!("dim_names: {:?}", xa.dim_names()); + println!("shape: {:?}\n", xa.shape()); + + // Coordinate selection uses the values attached to the named axis. + println!("=== Coordinate selection ===\n"); + + println!("xa.at(\"hour\", 2.0) - collapse to one observation"); + let at_two = xa.at("hour", 2.0); + println!(" shape {:?}, temp = {}\n", at_two.shape(), at_two.get(&[0])); + + println!("xa.between(\"hour\", 1.0, 3.0) - inclusive window"); + let window = xa.between("hour", 1.0, 3.0); + println!(" shape {:?}, first temp = {}\n", window.shape(), window.get(&[0, 0])); + + println!("xa.nearest(\"hour\", 1.8) - closest coordinate"); + let near = xa.nearest("hour", 1.8); + println!(" shape {:?}, temp = {}\n", near.shape(), near.get(&[0])); + + // Positional selection can also address an axis by name. + println!("=== Named positional selection ===\n"); + let first_two = xa.select(&[("hour", &(0..2))]); + println!("xa.select(&[(\"hour\", &(0..2))]): shape {:?}\n", first_two.shape()); + + // Transpose reorders data, dimension names, and coordinates together. + println!("=== Transpose ===\n"); + let t = xa.transpose(&["measurement", "hour"]).unwrap(); + println!("dims after transpose: {:?}, shape {:?}\n", t.dim_names(), t.shape()); + + // Axis-1 coordinate labels become Table column names. + println!("=== To Table ===\n"); + let mut named = xa.clone(); + named.assign_coords("measurement", arr_str32!(&["temp", "pressure", "humidity"])); + let table = named.to_table().unwrap(); + println!("{}", table); +} diff --git a/minarrow-py/Cargo.lock b/minarrow-py/Cargo.lock index b241cbf..66ca937 100644 --- a/minarrow-py/Cargo.lock +++ b/minarrow-py/Cargo.lock @@ -8,27 +8,12 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" -[[package]] -name = "cfg-if" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" - [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "libc" version = "0.2.186" @@ -41,15 +26,6 @@ version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "minarrow" version = "0.15.0" @@ -113,37 +89,32 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-build-config" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -151,9 +122,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -163,13 +134,12 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] @@ -183,12 +153,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - [[package]] name = "syn" version = "2.0.117" @@ -202,9 +166,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "thiserror" @@ -232,12 +196,6 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "vec64" version = "0.4.7" diff --git a/minarrow-py/Cargo.toml b/minarrow-py/Cargo.toml index 6aadd79..035aa7d 100644 --- a/minarrow-py/Cargo.toml +++ b/minarrow-py/Cargo.toml @@ -21,14 +21,27 @@ crate-type = ["cdylib", "rlib"] [dependencies] minarrow = { version = "0.15", path = ".." } minarrow-pyo3 = { version = "0.3", path = "../pyo3", optional = true, default-features = false } -pyo3 = { version = "0.23", features = ["abi3-py39"] } +pyo3 = { version = "0.29", features = ["abi3-py39"] } thiserror = "2" [features] -default = ["datetime", "large_string", "scalar_type", "value_type", "cube", "arrow_interop"] +default = ["datetime", "large_string", "scalar_type", "value_type", "cube", "arrow_interop", "ndarray"] extension-module = ["pyo3/extension-module"] arrow_interop = ["dep:minarrow-pyo3", "minarrow-pyo3/datetime"] simd = ["minarrow/simd"] +# N-dimensional tensor with DLPack capsule interchange. The capsule glue +# lives in minarrow-pyo3's ffi::dlpack, alongside the Arrow capsule glue. +ndarray = [ + "scalar_type", + "minarrow/ndarray", + "minarrow/dlpack", + "minarrow/views", + "minarrow/select", + "minarrow/chunked", + "minarrow/xarray", + "dep:minarrow-pyo3", + "minarrow-pyo3/ndarray", +] # Embeds a CPython interpreter so Rust can drive Python through `PyLiquid`. # Links libpython - use the default non `extension-module` link mode and @@ -66,8 +79,13 @@ name = "roundtrip_ml" path = "examples/roundtrip_ml.rs" required-features = ["embed"] +[[example]] +name = "roundtrip_tensor" +path = "examples/roundtrip_tensor.rs" +required-features = ["embed", "ndarray"] + [build-dependencies] -pyo3-build-config = "0.23" +pyo3-build-config = { version = "0.29", features = ["resolve-config"] } # Release profile is optimised, stripped, and carries no debug info. # diff --git a/minarrow-py/README.md b/minarrow-py/README.md index 6fd304e..8780109 100644 --- a/minarrow-py/README.md +++ b/minarrow-py/README.md @@ -33,7 +33,7 @@ Minarrow is useful when: * A small set of typed containers is preferable to a large object hierarchy * Higher-level computation is delegated to Polars, DuckDB, PyArrow or another execution engine -The Python API centres on four containers: +The columnar API centres on four containers: * `Array` for a single typed column * `Table` for a fixed set of equal-length columns @@ -42,6 +42,18 @@ The Python API centres on four containers: Together, they cover the common Arrow tabular data model. Nested `Struct` and `List` types are not currently supported. +The tensor API adds three focused containers: + +* `NdArray` for contiguous or zero-copy windowed f32/f64 data +* `ChunkedNdArray` for compatible `NdArray` pieces backed by Rust `SuperNdArray` +* `XArray` for named dimensions and coordinate-based selection over an `NdArray` + +They provide storage, indexing, selection, and DLPack interchange. `NdArray` +and `XArray` expose one tensor directly; each `ChunkedNdArray` chunk is a +separate `NdArray` DLPack producer with its own data pointer. Numerical +algorithms, statistics, and general tensor computation remain the job of +NumPy, PyTorch, JAX, or another execution library. + ## Why not just use PyArrow? PyArrow is an excellent Arrow implementation, but it solves a different problem. It brings the full Arrow C++ runtime into Python, together with a large type system, compute engine, file formats and dataset APIs. @@ -58,7 +70,7 @@ For example, a Rust service can use Lightstream to receive a network feed direct | Rust–Python data model | Native Minarrow types across both runtimes | Rust integration through Arrow interchange interfaces | | Host runtime | Rust, embedded in the application service | Arrow C++ runtime loaded into Python | | Buffer alignment | 64-byte aligned for Minarrow-backed buffers | 8-byte Arrow alignment guarantee | -| Python API | Four core containers for flat columnar data | Full Arrow type and container hierarchy | +| Python API | Focused columnar and tensor containers | Full Arrow type and container hierarchy | | Native SIMD use | Buffers are ready for aligned SIMD kernels | Consumers must inspect alignment or realign the data | | Runtime role | Application data model, interchange and composition | Arrow compute, storage and dataset platform | | Ecosystem integration | Arrow PyCapsule and direct runtime bridges | PyArrow APIs and Arrow interoperability | diff --git a/minarrow-py/docs/api.md b/minarrow-py/docs/api.md index e93f147..80a027a 100644 --- a/minarrow-py/docs/api.md +++ b/minarrow-py/docs/api.md @@ -13,6 +13,14 @@ documentation build environment. ::: minarrow.ChunkedTable +## N-dimensional data + +::: minarrow.NdArray + +::: minarrow.ChunkedNdArray + +::: minarrow.XArray + ## Data types ::: minarrow.DType diff --git a/minarrow-py/docs/index.md b/minarrow-py/docs/index.md index f628aa5..4ea07f9 100644 --- a/minarrow-py/docs/index.md +++ b/minarrow-py/docs/index.md @@ -32,7 +32,7 @@ --- - Array, Table, ChunkedArray and ChunkedTable cover flat columnar data. + Focused containers cover flat columnar and N-dimensional data. @@ -71,21 +71,35 @@ The same data model exists on both sides of the boundary. ## Overview -Minarrow provides four primary *Python* containers: +Minarrow provides four primary columnar *Python* containers: * `Array` for typed columnar data * `Table` for named collections of equal-length arrays * `ChunkedArray` for one logical column split across multiple chunks * `ChunkedTable` for a sequence of table batches -The *Python* package is backed by the [Minarrow Rust core](https://github.com/pbower/minarrow), which provides the underlying array types, schemas and aligned buffers. +For f32/f64 N-dimensional data it also includes: + +* `NdArray` for contiguous data and zero-copy selections +* `ChunkedNdArray` for batched array pieces - e.g., off disk or a network feed +* `XArray` for named axes and coordinate selection + +These are intended as foundational data containers and interchange objects, +rather than a direct replacement for a numerical or statistical runtime at this layer. +Use their DLPack support to pass data to NumPy, PyTorch, JAX, or another +compute library. `ChunkedNdArray` exposes one DLPack-producing `NdArray` per +chunk rather than consolidating its allocations implicitly. + +The *Python* package is backed by [Minarrow Rust](https://github.com/pbower/minarrow), which handles the underlying array +operations, schemas and 64-byte aligned buffers (for SIMD), at high levels of performance. Minarrow is intended for: -* Passing columnar data between *Rust* and *Python* +* Tabular and NdArray application landing zone for high-performance use cases. +* Data interchange between *Rust* and *Python* * Building *Python* extensions that operate on *Arrow*-compatible data * Feeding *Polars*, *DuckDB*, *PyArrow* and other *Arrow*-aware libraries -* Running SIMD-oriented native kernels over guaranteed aligned buffers +* Running SIMD-oriented native statistical kernels. * Keeping binary size and dependency footprint small ## Installation @@ -179,7 +193,7 @@ Use it to construct or receive data, operate on it through native extensions, an | *pandas* workflows | *pandas* with an *Arrow*-compatible path | | Custom native computation | *Rust*, C or *C++* through *Arrow* capsules | -The [Rust crate](https://docs.rs/minarrow/latest/minarrow/) does provide a minimal set of tabular data-processing capabilities but without full parallelism, which serves as a useful basis for foundational operations in that environment. It compiles in less than 2 seconds with the standard feature set to help you stay productive, and your project lightweight. +The [Rust crate](https://docs.rs/minarrow/latest/minarrow/) does provide a minimal set of tabular data-processing capabilities but without full parallelism, which serves as a useful basis for foundational operations in that environment. It compiles in less than 2 seconds with the standard feature set to help you stay productive, and your project lightweight. See [Ecosystem interoperability](interop.md) for supported conversion paths and ownership behaviour. diff --git a/minarrow-py/docs/interop.md b/minarrow-py/docs/interop.md index 5cf2b61..c61464f 100644 --- a/minarrow-py/docs/interop.md +++ b/minarrow-py/docs/interop.md @@ -40,6 +40,22 @@ Whether an integration remains zero-copy depends on the receiving library, data The PyCapsule interface is part of the Apache Arrow interoperability specifications. Minarrow implements those interfaces independently and does not require PyArrow to own or manage its buffers. +## DLPack tensor interchange + +Minarrow supports DLPack for cross-language tensor interchange. + +This enables you to land data for e.g., in Rust, when using the Python bridge, +get the data into PyTorch for Deep Learning, AI or other similar use case. + +`NdArray` and `XArray` implement `__dlpack__` and `__dlpack_device__` for CPU +tensor interchange. Behaviour-wise, their owned arrays and zero-copy selections +retain shape and element strides, including Minarrow's column-major layout. + +`ChunkedNdArray` yields individual arrays. Its `chunks` are `NdArray` objects, +and each chunk implements DLPack with its own shape, strides, and data pointer. +`to_numpy()` consequently returns one NumPy array per chunk. +Call `to_ndarray()` when one consolidated allocation is required. + ## Supported integrations Minarrow provides named conversion methods for commonly used libraries. diff --git a/minarrow-py/examples/pyliquid_ecosystem.rs b/minarrow-py/examples/pyliquid_ecosystem.rs index 78bb51c..0e79c15 100644 --- a/minarrow-py/examples/pyliquid_ecosystem.rs +++ b/minarrow-py/examples/pyliquid_ecosystem.rs @@ -269,14 +269,14 @@ fn bridge_costs(small: &Table) -> PyResult<()> { let t = Instant::now(); for _ in 0..G { - Python::with_gil(|_py| {}); + Python::attach(|_py| {}); } println!(" GIL acquire and release : {:>11.3?}", t.elapsed() / G); let large = build_large(1_000_000); - Python::with_gil(|py| -> PyResult<()> { + Python::attach(|py| -> PyResult<()> { const N: u32 = 50; println!(); diff --git a/minarrow-py/examples/pyliquid_lifetime.rs b/minarrow-py/examples/pyliquid_lifetime.rs index da13672..e9cbd30 100644 --- a/minarrow-py/examples/pyliquid_lifetime.rs +++ b/minarrow-py/examples/pyliquid_lifetime.rs @@ -40,7 +40,7 @@ fn main() -> PyResult<()> { // Build the input table in Rust. let ids: Vec64 = (0..ROWS).collect(); let input = Table::new("input".to_string(), Some(vec![fa_i64!("x", @vec64 ids)])); - let input_address = input.cols[0].array.try_i64_ref().unwrap().data.as_ptr() as usize; + let input_address = input.cols[0].array.num().i64().data.as_ptr() as usize; // Pass the table into Python, multiply the column in Polars, and return the // resulting frame. That dataframe is the object imported back into Minarrow. @@ -49,11 +49,11 @@ fn main() -> PyResult<()> { scope.set_item("table", obj)?; py.run( cr#" - import polars as pl +import polars as pl - df = table.to_polars() - result = df.with_columns((pl.col("x") * 7).alias("x")) - "#, +df = table.to_polars() +result = df.with_columns((pl.col("x") * 7).alias("x")) +"#, Some(&scope), Some(&scope), )?; @@ -66,7 +66,7 @@ fn main() -> PyResult<()> { Value::Table(table) => (*table).clone(), other => panic!("expected a table, got {other:?}"), }; - let result_address = result.cols[0].array.try_i64_ref().unwrap().data.as_ptr() as usize; + let result_address = result.cols[0].array.num().i64().data.as_ptr() as usize; println!("Received the Polars result into Rust."); println!(" Rust input buffer : {input_address:#x}"); println!(" Polars result buffer : {result_address:#x}"); @@ -85,7 +85,7 @@ fn main() -> PyResult<()> { sleep(Duration::from_secs(1)); // Use the Polars result after Python is gone. - let sum: i64 = result.cols[0].array.clone().num().i64().unwrap().data.iter().sum(); + let sum: i64 = result.cols[0].array.clone().num().i64().data.iter().sum(); let expected: i64 = (0..ROWS).map(|value| value * 7).sum(); println!( "After shutdown: rows = {}, sum = {sum} (expected {expected})", diff --git a/minarrow-py/examples/roundtrip_ml.rs b/minarrow-py/examples/roundtrip_ml.rs index d10f02d..d2ed03a 100644 --- a/minarrow-py/examples/roundtrip_ml.rs +++ b/minarrow-py/examples/roundtrip_ml.rs @@ -53,27 +53,27 @@ fn main() -> PyResult<()> { scope.set_item("table", obj)?; py.run( cr#" - import polars as pl - from sklearn.ensemble import RandomForestClassifier - from sklearn.model_selection import train_test_split +import polars as pl +from sklearn.ensemble import RandomForestClassifier +from sklearn.model_selection import train_test_split - df = table.to_polars() - features = ["x0", "x1", "x2", "x3"] - X = df.select(features).to_numpy() - y = df["label"].to_numpy() +df = table.to_polars() +features = ["x0", "x1", "x2", "x3"] +X = df.select(features).to_numpy() +y = df["label"].to_numpy() - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0) - model = RandomForestClassifier(n_estimators=100, random_state=0) - model.fit(X_train, y_train) - predicted = model.predict(X_test) +X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0) +model = RandomForestClassifier(n_estimators=100, random_state=0) +model.fit(X_train, y_train) +predicted = model.predict(X_test) - result = pl.DataFrame( - { - "actual": y_test.astype("int64"), - "predicted": predicted.astype("int64"), - } - ) - "#, +result = pl.DataFrame( + { + "actual": y_test.astype("int64"), + "predicted": predicted.astype("int64"), + } +) +"#, Some(&scope), Some(&scope), )?; @@ -85,8 +85,8 @@ fn main() -> PyResult<()> { let Value::Table(table) = value else { panic!("expected a table, got {value:?}"); }; - let actual = table.cols[0].array.try_i64_ref().unwrap(); - let predicted = table.cols[1].array.try_i64_ref().unwrap(); + let actual = table.cols[0].array.num().i64(); + let predicted = table.cols[1].array.num().i64(); let rows = actual.len(); let correct = (0..rows).filter(|&i| actual.data[i] == predicted.data[i]).count(); diff --git a/minarrow-py/examples/roundtrip_tensor.rs b/minarrow-py/examples/roundtrip_tensor.rs new file mode 100644 index 0000000..a75e6f9 --- /dev/null +++ b/minarrow-py/examples/roundtrip_tensor.rs @@ -0,0 +1,104 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Round-trips a tensor through PyTorch over DLPack. +//! +//! An `NdArray` of features is built in Rust and passed into Python, where +//! `torch.from_dlpack` picks up the buffer zero-copy. The closure +//! standardises each feature column, scores every row against a weight +//! vector, and returns the score tensor. That tensor crosses back over the +//! DLPack capsule protocol and Rust reads the results from Minarrow's own +//! 64-byte aligned memory. +//! +//! ## Running +//! Needs a venv with `torch`. `PYO3_PYTHON` must be an absolute path to it. +//! ```bash +//! cd minarrow-py +//! PYO3_PYTHON=$PWD/../pyo3/.venv/bin/python \ +//! PYTHONHOME=/usr \ +//! PYTHONPATH=$PWD/../pyo3/.venv/lib/python3.12/site-packages \ +//! LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu \ +//! cargo run --release --example roundtrip_tensor --features embed,ndarray +//! ``` + +use minarrow::{NdArray, Value, Vec64}; +use minarrow_py::PyLiquid; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +fn main() -> PyResult<()> { + let rt = PyLiquid::start().preimport(["torch"])?; + + let features = build_features(1024, 8); + println!( + "Built a {:?} tensor in Rust at {:p}.", + features.shape(), + features.as_slice().as_ptr() + ); + + // Hand the tensor to PyTorch, standardise the feature columns, and + // score every row against a fixed weight vector. + let value = rt.with_python(&features, |py, obj| { + let scope = PyDict::new(py); + scope.set_item("tensor", obj)?; + py.run( + cr#" +import torch + +x = torch.from_dlpack(tensor) +mean = x.mean(dim=0, keepdim=True) +std = x.std(dim=0, keepdim=True) +z = (x - mean) / std +weights = torch.linspace(1.0, 2.0, x.shape[1], dtype=torch.float64) +result = (z * weights).sum(dim=1) +"#, + Some(&scope), + Some(&scope), + )?; + scope + .get_item("result")? + .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err("result not set")) + })?; + + let Value::NdArray(scores) = value else { + panic!("expected a tensor, got {value:?}"); + }; + println!("Scored the rows in PyTorch and read them back in Rust:"); + println!(" shape : {:?}", scores.shape()); + println!(" first three : {:?}", &scores.as_slice()[..3]); + println!( + " mean score : {:.4}", + scores.as_slice().iter().sum::() / scores.len() as f64 + ); + + Ok(()) +} + +/// Builds a feature tensor from a deterministic xorshift generator, laid +/// out column-major as `[rows, cols]`. +fn build_features(rows: usize, cols: usize) -> NdArray { + let mut state: u64 = 0x9E37_79B9_7F4A_7C15; + let mut draw = move || -> f64 { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + ((state >> 11) as f64 / (1u64 << 53) as f64) * 2.0 - 1.0 + }; + + let mut flat = Vec64::with_capacity(rows * cols); + for _ in 0..rows * cols { + flat.push(draw()); + } + NdArray::from_vec64(flat, &[rows, cols]) +} diff --git a/minarrow-py/src/array.rs b/minarrow-py/src/array.rs index 82debaf..cd3c8b3 100644 --- a/minarrow-py/src/array.rs +++ b/minarrow-py/src/array.rs @@ -188,12 +188,12 @@ impl PyArrayInner { &self, py: Python<'_>, key: &Bound<'_, PyAny>, - wrap_window: impl Fn(ArrayV) -> PyResult, - ) -> PyResult { + wrap_window: impl Fn(ArrayV) -> PyResult>, + ) -> PyResult> { let view: ArrayV = self.into(); let len = view.len(); - if let Ok(slice) = key.downcast::() { + if let Ok(slice) = key.cast::() { let ind = slice.indices(len as isize)?; let windowed = if ind.step == 1 { view.r((ind.start as usize)..(ind.stop as usize)) @@ -538,7 +538,7 @@ impl PyArray { self.0.repr() } - fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult { + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { self.0.get_item(py, key, |window| { Ok(Py::new(py, PyArray::from(window))?.into_any()) }) @@ -599,8 +599,8 @@ impl PyArray { fn __arrow_c_array__( &self, py: Python<'_>, - requested_schema: Option, - ) -> PyResult<(PyObject, PyObject)> { + requested_schema: Option>, + ) -> PyResult<(Py, Py)> { let _ = requested_schema; match &self.0 { PyArrayInner::Array(array) => { diff --git a/minarrow-py/src/arrow_type.rs b/minarrow-py/src/arrow_type.rs index a78bf90..994294d 100644 --- a/minarrow-py/src/arrow_type.rs +++ b/minarrow-py/src/arrow_type.rs @@ -27,7 +27,7 @@ use pyo3::prelude::*; /// The unit of a temporal type. Mirrors `minarrow::TimeUnit`. #[cfg(feature = "datetime")] -#[pyclass(eq, eq_int, name = "TimeUnit", module = "minarrow")] +#[pyclass(from_py_object, eq, eq_int, name = "TimeUnit", module = "minarrow")] #[derive(Clone, Copy, PartialEq)] pub enum PyTimeUnit { Seconds, @@ -65,7 +65,7 @@ impl From for TimeUnit { /// The unit of an interval type. Mirrors `minarrow::IntervalUnit`. #[cfg(feature = "datetime")] -#[pyclass(eq, eq_int, name = "IntervalUnit", module = "minarrow")] +#[pyclass(from_py_object, eq, eq_int, name = "IntervalUnit", module = "minarrow")] #[derive(Clone, Copy, PartialEq)] pub enum PyIntervalUnit { YearMonth, @@ -97,7 +97,7 @@ impl From for IntervalUnit { /// The dictionary key width of a categorical type. Mirrors /// `minarrow::CategoricalIndexType` under its feature gates. -#[pyclass(eq, eq_int, name = "CategoricalIndexType", module = "minarrow")] +#[pyclass(from_py_object, eq, eq_int, name = "CategoricalIndexType", module = "minarrow")] #[derive(Clone, Copy, PartialEq)] pub enum PyCategoricalIndexType { #[cfg(feature = "default_categorical_8")] @@ -144,7 +144,7 @@ impl From for CategoricalIndexType { /// feature gates. Construct it for a `Field`, or read it from `Array.arrow_type`. /// pyo3 makes each variant callable, so a non-parametric type is built with a /// call: `ArrowType.Int64()`. -#[pyclass(eq, name = "ArrowType", module = "minarrow")] +#[pyclass(from_py_object, eq, name = "ArrowType", module = "minarrow")] #[derive(Clone, PartialEq)] pub enum PyArrowType { Null(), diff --git a/minarrow-py/src/chunked_array.rs b/minarrow-py/src/chunked_array.rs index f324a10..f902d22 100644 --- a/minarrow-py/src/chunked_array.rs +++ b/minarrow-py/src/chunked_array.rs @@ -147,8 +147,8 @@ impl PyChunkedArray { fn __arrow_c_stream__( &self, py: Python<'_>, - requested_schema: Option, - ) -> PyResult { + requested_schema: Option>, + ) -> PyResult> { let _ = requested_schema; to_py::super_array_to_stream_capsule(&self.0, py) } diff --git a/minarrow-py/src/chunked_ndarray.rs b/minarrow-py/src/chunked_ndarray.rs new file mode 100644 index 0000000..75c4c01 --- /dev/null +++ b/minarrow-py/src/chunked_ndarray.rs @@ -0,0 +1,313 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! `ChunkedNdArray` - Python's name for minarrow's `SuperNdArray`. + +use std::sync::Arc; + +use minarrow::{AxisSelection, Consolidate, NdArray, SuperNdArray, SuperNdArrayV}; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +use crate::ndarray::{AxisKey, PyNdArray, PyNdArrayInner, axis_keys, selectors}; + +pub enum PyChunkedNdArrayInner { + F32(Arc>), + F64(Arc>), + F32View(Arc>), + F64View(Arc>), +} + +/// An ordered set of compatible NdArray pieces, backed by SuperNdArray. +#[pyclass(name = "ChunkedNdArray", module = "minarrow")] +pub struct PyChunkedNdArray(pub PyChunkedNdArrayInner); + +fn validate_chunks(chunks: &[NdArray]) -> PyResult<()> { + let Some(first) = chunks.first() else { + return Ok(()); + }; + if first.ndim() == 0 { + return Err(PyValueError::new_err( + "ChunkedNdArray pieces require an axis 0", + )); + } + for (index, chunk) in chunks.iter().enumerate().skip(1) { + if chunk.ndim() != first.ndim() || chunk.shape()[1..] != first.shape()[1..] { + return Err(PyValueError::new_err(format!( + "chunk {index} has shape {:?}, expected rank {} and trailing shape {:?}", + chunk.shape(), + first.ndim(), + &first.shape()[1..] + ))); + } + } + Ok(()) +} + +impl From> for PyChunkedNdArray { + fn from(array: SuperNdArray) -> Self { + Self(PyChunkedNdArrayInner::F32(Arc::new(array))) + } +} + +impl From> for PyChunkedNdArray { + fn from(array: SuperNdArray) -> Self { + Self(PyChunkedNdArrayInner::F64(Arc::new(array))) + } +} + +impl From> for PyChunkedNdArray { + fn from(view: SuperNdArrayV) -> Self { + Self(PyChunkedNdArrayInner::F32View(Arc::new(view))) + } +} + +impl From> for PyChunkedNdArray { + fn from(view: SuperNdArrayV) -> Self { + Self(PyChunkedNdArrayInner::F64View(Arc::new(view))) + } +} + +#[pymethods] +impl PyChunkedNdArray { + /// Construct from compatible NdArray pieces. Empty input defaults to + /// float64 unless `dtype="float32"` is supplied. + #[new] + #[pyo3(signature = (chunks, name="", dtype=None))] + fn new(chunks: Vec>, name: &str, dtype: Option<&str>) -> PyResult { + let dtype = dtype.unwrap_or_else(|| { + chunks + .first() + .map_or("float64", |chunk| match &chunk.borrow().0 { + PyNdArrayInner::F32(_) | PyNdArrayInner::F32View(_) => "float32", + PyNdArrayInner::F64(_) | PyNdArrayInner::F64View(_) => "float64", + }) + }); + match dtype { + "float32" | "f32" => { + let arrays = chunks + .iter() + .map(|chunk| match &chunk.borrow().0 { + PyNdArrayInner::F32(a) => Ok((**a).clone()), + PyNdArrayInner::F32View(v) => Ok(v.to_ndarray()), + _ => Err(PyTypeError::new_err( + "all ChunkedNdArray pieces must have dtype float32", + )), + }) + .collect::>>>()?; + validate_chunks(&arrays)?; + Ok(SuperNdArray::from_batches(arrays, name).into()) + } + "float64" | "f64" => { + let arrays = chunks + .iter() + .map(|chunk| match &chunk.borrow().0 { + PyNdArrayInner::F64(a) => Ok((**a).clone()), + PyNdArrayInner::F64View(v) => Ok(v.to_ndarray()), + _ => Err(PyTypeError::new_err( + "all ChunkedNdArray pieces must have dtype float64", + )), + }) + .collect::>>>()?; + validate_chunks(&arrays)?; + Ok(SuperNdArray::from_batches(arrays, name).into()) + } + other => Err(PyValueError::new_err(format!( + "dtype must be 'float32' or 'float64', got '{other}'" + ))), + } + } + + #[getter] + fn shape<'py>(&self, py: Python<'py>) -> PyResult> { + PyTuple::new(py, self.shape_vec()) + } + + #[getter] + fn ndim(&self) -> usize { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.ndim(), + PyChunkedNdArrayInner::F64(a) => a.ndim(), + PyChunkedNdArrayInner::F32View(v) => v.ndim(), + PyChunkedNdArrayInner::F64View(v) => v.ndim(), + } + } + + #[getter] + fn dtype(&self) -> &'static str { + match &self.0 { + PyChunkedNdArrayInner::F32(_) | PyChunkedNdArrayInner::F32View(_) => "float32", + PyChunkedNdArrayInner::F64(_) | PyChunkedNdArrayInner::F64View(_) => "float64", + } + } + + #[getter] + fn size(&self) -> usize { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.len(), + PyChunkedNdArrayInner::F64(a) => a.len(), + PyChunkedNdArrayInner::F32View(v) => v.len(), + PyChunkedNdArrayInner::F64View(v) => v.len(), + } + } + + #[getter] + fn n_chunks(&self) -> usize { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.n_batches(), + PyChunkedNdArrayInner::F64(a) => a.n_batches(), + PyChunkedNdArrayInner::F32View(v) => v.n_slices(), + PyChunkedNdArrayInner::F64View(v) => v.n_slices(), + } + } + + #[getter] + fn is_view(&self) -> bool { + matches!( + &self.0, + PyChunkedNdArrayInner::F32View(_) | PyChunkedNdArrayInner::F64View(_) + ) + } + + #[getter] + fn name(&self) -> Option { + let name = match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.name.as_str(), + PyChunkedNdArrayInner::F64(a) => a.name.as_str(), + PyChunkedNdArrayInner::F32View(v) => v.name(), + PyChunkedNdArrayInner::F64View(v) => v.name(), + }; + (!name.is_empty()).then(|| name.to_string()) + } + + /// The constituent pieces in order. A window returns its zero-copy + /// NdArray windows rather than materialising them. + #[getter] + fn chunks(&self) -> Vec { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => { + a.batches().iter().cloned().map(PyNdArray::from).collect() + } + PyChunkedNdArrayInner::F64(a) => { + a.batches().iter().cloned().map(PyNdArray::from).collect() + } + PyChunkedNdArrayInner::F32View(v) => { + v.slices.iter().cloned().map(PyNdArray::from).collect() + } + PyChunkedNdArrayInner::F64View(v) => { + v.slices.iter().cloned().map(PyNdArray::from).collect() + } + } + } + + /// One constituent piece by position, or `None` when out of range. + fn chunk(&self, index: usize) -> Option { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.batch(index).cloned().map(PyNdArray::from), + PyChunkedNdArrayInner::F64(a) => a.batch(index).cloned().map(PyNdArray::from), + PyChunkedNdArrayInner::F32View(v) => v.slices.get(index).cloned().map(PyNdArray::from), + PyChunkedNdArrayInner::F64View(v) => v.slices.get(index).cloned().map(PyNdArray::from), + } + } + + fn __len__(&self) -> usize { + self.shape_vec()[0] + } + + fn __repr__(&self) -> String { + format!( + "ChunkedNdArray(shape={:?}, dtype={}, chunks={})", + self.shape_vec(), + self.dtype(), + self.n_chunks() + ) + } + + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { + let shape = self.shape_vec(); + let keys = axis_keys(key, &shape)?; + if keys.iter().all(|key| matches!(key, AxisKey::Index(_))) { + let indices: Vec = keys + .iter() + .map(|key| match key { + AxisKey::Index(index) => *index, + AxisKey::Range(_) => unreachable!(), + }) + .collect(); + let value = match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.get(&indices) as f64, + PyChunkedNdArrayInner::F64(a) => a.get(&indices), + PyChunkedNdArrayInner::F32View(v) => v.get(&indices) as f64, + PyChunkedNdArrayInner::F64View(v) => v.get(&indices), + }; + return value.into_py_any(py); + } + + if let AxisKey::Index(index) = keys[0] { + let trailing = selectors(&keys[1..]); + let array = match &self.0 { + PyChunkedNdArrayInner::F32(a) => PyNdArray::from(a.obs(index).slice(&trailing)), + PyChunkedNdArrayInner::F64(a) => PyNdArray::from(a.obs(index).slice(&trailing)), + PyChunkedNdArrayInner::F32View(v) => PyNdArray::from(v.obs(index).slice(&trailing)), + PyChunkedNdArrayInner::F64View(v) => PyNdArray::from(v.obs(index).slice(&trailing)), + }; + return Ok(Py::new(py, array)?.into_any()); + } + + let selection = selectors(&keys); + let view = match &self.0 { + PyChunkedNdArrayInner::F32(a) => PyChunkedNdArray::from(a.s(&selection)), + PyChunkedNdArrayInner::F64(a) => PyChunkedNdArray::from(a.s(&selection)), + PyChunkedNdArrayInner::F32View(v) => PyChunkedNdArray::from(v.s(&selection)), + PyChunkedNdArrayInner::F64View(v) => PyChunkedNdArray::from(v.s(&selection)), + }; + Ok(Py::new(py, view)?.into_any()) + } + + /// Materialise to one compact column-major NdArray. + fn to_ndarray(&self) -> PyNdArray { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => PyNdArray::from((**a).clone().consolidate()), + PyChunkedNdArrayInner::F64(a) => PyNdArray::from((**a).clone().consolidate()), + PyChunkedNdArrayInner::F32View(v) => PyNdArray::from((**v).clone().consolidate()), + PyChunkedNdArrayInner::F64View(v) => PyNdArray::from((**v).clone().consolidate()), + } + } + + /// Hand each chunk to NumPy through its own DLPack producer. The returned + /// list preserves one tensor and one data pointer per chunk. + fn to_numpy(&self, py: Python<'_>) -> PyResult>> { + let numpy = py.import("numpy")?; + self.chunks() + .into_iter() + .map(|chunk| { + let chunk = Py::new(py, chunk)?; + Ok(numpy.call_method1("from_dlpack", (chunk,))?.unbind()) + }) + .collect() + } +} + +impl PyChunkedNdArray { + fn shape_vec(&self) -> Vec { + match &self.0 { + PyChunkedNdArrayInner::F32(a) => a.shape(), + PyChunkedNdArrayInner::F64(a) => a.shape(), + PyChunkedNdArrayInner::F32View(v) => v.shape(), + PyChunkedNdArrayInner::F64View(v) => v.shape(), + } + } +} diff --git a/minarrow-py/src/chunked_table.rs b/minarrow-py/src/chunked_table.rs index 49ee024..ad379ab 100644 --- a/minarrow-py/src/chunked_table.rs +++ b/minarrow-py/src/chunked_table.rs @@ -164,8 +164,8 @@ impl PyChunkedTable { fn __arrow_c_stream__( &self, py: Python<'_>, - requested_schema: Option, - ) -> PyResult { + requested_schema: Option>, + ) -> PyResult> { let _ = requested_schema; to_py::super_table_to_stream_capsule(&self.0, py) } diff --git a/minarrow-py/src/convert.rs b/minarrow-py/src/convert.rs index fc570d0..f772395 100644 --- a/minarrow-py/src/convert.rs +++ b/minarrow-py/src/convert.rs @@ -347,7 +347,7 @@ pub fn resolve_index(i: isize, len: usize) -> PyResult { /// Temporal values surface as their raw integer. Faithful /// `datetime.date`/`time`/`datetime` coercion needs the logical unit and /// timezone carried on the `Field` and lands with the temporal family. -pub fn scalar_to_py(py: Python<'_>, scalar: Scalar) -> PyResult { +pub fn scalar_to_py(py: Python<'_>, scalar: Scalar) -> PyResult> { match scalar { Scalar::Null => Ok(py.None()), Scalar::Boolean(v) => v.into_py_any(py), diff --git a/minarrow-py/src/dtype.rs b/minarrow-py/src/dtype.rs index cfd2911..41fa6d4 100644 --- a/minarrow-py/src/dtype.rs +++ b/minarrow-py/src/dtype.rs @@ -23,7 +23,7 @@ use minarrow::ffi::arrow_dtype::{ArrowType, CategoricalIndexType}; use pyo3::prelude::*; /// minarrow's `Array` enum grouping. -#[pyclass(eq, eq_int, name = "TypeClass", module = "minarrow")] +#[pyclass(from_py_object, eq, eq_int, name = "TypeClass", module = "minarrow")] #[derive(Clone, Copy, PartialEq)] pub enum TypeClass { Numeric, @@ -58,7 +58,7 @@ impl TypeClass { } /// The concrete minarrow array type. -#[pyclass(eq, eq_int, name = "DType", module = "minarrow")] +#[pyclass(from_py_object, eq, eq_int, name = "DType", module = "minarrow")] #[derive(Clone, Copy, PartialEq)] pub enum DType { Integer, diff --git a/minarrow-py/src/field.rs b/minarrow-py/src/field.rs index db4a133..01f4cde 100644 --- a/minarrow-py/src/field.rs +++ b/minarrow-py/src/field.rs @@ -29,7 +29,7 @@ use crate::convert::resolve_index; use crate::dtype::{dtype_from_arrow, DType}; /// A named column descriptor holding name, Arrow type, nullability, and metadata. -#[pyclass(name = "Field", module = "minarrow")] +#[pyclass(from_py_object, name = "Field", module = "minarrow")] #[derive(Clone)] pub struct PyField(pub Field); @@ -93,7 +93,7 @@ impl PyField { } /// An ordered list of fields plus table-level metadata. -#[pyclass(name = "Schema", module = "minarrow")] +#[pyclass(from_py_object, name = "Schema", module = "minarrow")] #[derive(Clone)] pub struct PySchema(pub Schema); diff --git a/minarrow-py/src/lib.rs b/minarrow-py/src/lib.rs index ce99c90..3572f01 100644 --- a/minarrow-py/src/lib.rs +++ b/minarrow-py/src/lib.rs @@ -21,13 +21,19 @@ mod array; mod arrow_type; mod chunked_array; +#[cfg(feature = "ndarray")] +mod chunked_ndarray; mod chunked_table; mod convert; mod dtype; mod field; +#[cfg(feature = "ndarray")] +mod ndarray; #[cfg(feature = "embed")] mod pyliquid; mod table; +#[cfg(feature = "ndarray")] +mod xarray; use pyo3::prelude::*; @@ -39,13 +45,19 @@ pub use arrow_type::{PyArrowType, PyCategoricalIndexType}; #[cfg(feature = "datetime")] pub use arrow_type::{PyIntervalUnit, PyTimeUnit}; pub use chunked_array::PyChunkedArray; +#[cfg(feature = "ndarray")] +pub use chunked_ndarray::{PyChunkedNdArray, PyChunkedNdArrayInner}; pub use chunked_table::PyChunkedTable; pub use convert::{build_array, resolve_index, scalar_to_py}; pub use dtype::{dtype_from_arrow, width_from_arrow, DType, TypeClass}; pub use field::{PyField, PySchema}; +#[cfg(feature = "ndarray")] +pub use ndarray::{PyNdArray, PyNdArrayInner}; #[cfg(feature = "embed")] pub use pyliquid::{PyInput, PyLiquid}; pub use table::{build_table, PyTableInner}; +#[cfg(feature = "ndarray")] +pub use xarray::{PyXArray, PyXArrayInner}; /// Registers the `minarrow` module in the embedded interpreter's inittab so /// `import minarrow` resolves. Kept here in the crate root, where the @@ -67,6 +79,12 @@ fn minarrow_py(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + #[cfg(feature = "ndarray")] + { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + } m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/minarrow-py/src/ndarray.rs b/minarrow-py/src/ndarray.rs new file mode 100644 index 0000000..3c430d9 --- /dev/null +++ b/minarrow-py/src/ndarray.rs @@ -0,0 +1,389 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! `NdArray` - the native Python tensor object over minarrow. +//! +//! One Python type holds an f32 or f64 `NdArray` or `NdArrayV`. Slicing +//! keeps the same public type and stores a zero-copy view internally. The +//! DLPack capsule protocol (`__dlpack__` and `__dlpack_device__`) hands the +//! tensor to NumPy, PyTorch, JAX, TensorFlow, and CuPy. Transfers share data +//! when the consumer, ownership, layout, and protocol permit it; otherwise +//! they copy or report an unsupported request. `from_dlpack` accepts supported +//! CPU f32/f64 producers. Named bridge methods use the same capsule path and +//! import the target library at call time. The capsule glue lives in +//! minarrow-pyo3's `ffi::dlpack`, alongside the shared Arrow capsule glue. + +use minarrow::traits::selection::DataSelector; +use minarrow::{NdArray, NdArrayV, Table}; +use minarrow_pyo3::ffi::dlpack::{export_dlpack, import_dlpack}; +pub use minarrow_pyo3::ffi::dlpack::PyNdArrayInner; + +use crate::table::{PyTable, PyTableInner}; +use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PySlice, PyTuple}; +use pyo3::IntoPyObjectExt; + +use crate::convert::resolve_index; + +/// Dispatch across dtype and owned/view storage. +macro_rules! dispatch { + ($inner:expr, |$a:ident| $body:expr) => { + match $inner { + PyNdArrayInner::F32($a) => $body, + PyNdArrayInner::F64($a) => $body, + PyNdArrayInner::F32View($a) => $body, + PyNdArrayInner::F64View($a) => $body, + } + }; +} + +pub(crate) enum AxisKey { + Index(usize), + Range(std::ops::Range), +} + +pub(crate) fn axis_keys(key: &Bound<'_, PyAny>, shape: &[usize]) -> PyResult> { + let items: Vec> = if let Ok(tuple) = key.cast::() { + tuple.iter().collect() + } else { + vec![key.clone()] + }; + if items.len() > shape.len() { + return Err(PyIndexError::new_err(format!( + "too many indices for {}-dimensional array", + shape.len() + ))); + } + + let mut keys = Vec::with_capacity(shape.len()); + for (axis, item) in items.iter().enumerate() { + if let Ok(slice) = item.cast::() { + let indices = slice.indices(shape[axis] as isize)?; + if indices.step != 1 { + return Err(PyValueError::new_err( + "NdArray slices currently require a step of 1", + )); + } + keys.push(AxisKey::Range( + (indices.start as usize)..(indices.stop as usize), + )); + } else if let Ok(index) = item.extract::() { + keys.push(AxisKey::Index(resolve_index(index, shape[axis])?)); + } else { + return Err(PyTypeError::new_err( + "each index must be an int or a slice", + )); + } + } + for &size in &shape[keys.len()..] { + keys.push(AxisKey::Range(0..size)); + } + Ok(keys) +} + +pub(crate) fn selectors(keys: &[AxisKey]) -> Vec<&dyn DataSelector> { + keys.iter() + .map(|key| match key { + AxisKey::Index(index) => index as &dyn DataSelector, + AxisKey::Range(range) => range as &dyn DataSelector, + }) + .collect() +} + +/// N-dimensional f32/f64 tensor with DLPack interchange and zero-copy sharing +/// where ownership, layout, and protocol permit it. +#[pyclass(name = "NdArray", module = "minarrow")] +pub struct PyNdArray(pub PyNdArrayInner); + +impl From> for PyNdArray { + fn from(ndarray: NdArray) -> Self { + PyNdArray(PyNdArrayInner::from(ndarray)) + } +} + +impl From> for PyNdArray { + fn from(ndarray: NdArray) -> Self { + PyNdArray(PyNdArrayInner::from(ndarray)) + } +} + +impl From> for PyNdArray { + fn from(view: NdArrayV) -> Self { + PyNdArray(PyNdArrayInner::from(view)) + } +} + +impl From> for PyNdArray { + fn from(view: NdArrayV) -> Self { + PyNdArray(PyNdArrayInner::from(view)) + } +} + +#[pymethods] +impl PyNdArray { + /// Build from a flat sequence of numbers, shaped column-major. + /// `shape` defaults to one dimension covering the whole sequence. + /// `shape=[]` creates a rank-zero array from exactly one value and is + /// distinct from `shape=[0]`, which creates an empty rank-one array. + #[new] + #[pyo3(signature = (data, shape=None, dtype="float64"))] + fn new(data: Vec, shape: Option>, dtype: &str) -> PyResult { + let shape = shape.unwrap_or_else(|| vec![data.len()]); + let expected: usize = shape.iter().product(); + if expected != data.len() { + return Err(PyValueError::new_err(format!( + "shape {:?} needs {} elements, data has {}", + shape, + expected, + data.len() + ))); + } + match dtype { + "float64" | "f64" => Ok(PyNdArray::from(NdArray::from_slice(&data, &shape))), + "float32" | "f32" => { + let data32: Vec = data.iter().map(|&v| v as f32).collect(); + Ok(PyNdArray::from(NdArray::from_slice(&data32, &shape))) + } + other => Err(PyValueError::new_err(format!( + "dtype must be 'float32' or 'float64', got '{}'", + other + ))), + } + } + + /// Dimension sizes. + #[getter] + fn shape<'py>(&self, py: Python<'py>) -> PyResult> { + dispatch!(&self.0, |a| PyTuple::new(py, a.shape())) + } + + /// Element strides per dimension. + #[getter] + fn strides<'py>(&self, py: Python<'py>) -> PyResult> { + dispatch!(&self.0, |a| PyTuple::new(py, a.strides())) + } + + /// Number of dimensions. + #[getter] + fn ndim(&self) -> usize { + dispatch!(&self.0, |a| a.ndim()) + } + + /// Element type name, `float32` or `float64`. + #[getter] + fn dtype(&self) -> &'static str { + match &self.0 { + PyNdArrayInner::F32(_) | PyNdArrayInner::F32View(_) => "float32", + PyNdArrayInner::F64(_) | PyNdArrayInner::F64View(_) => "float64", + } + } + + /// Whether this object is a zero-copy window over another NdArray. + #[getter] + fn is_view(&self) -> bool { + matches!( + &self.0, + PyNdArrayInner::F32View(_) | PyNdArrayInner::F64View(_) + ) + } + + /// Total element count. + #[getter] + fn size(&self) -> usize { + dispatch!(&self.0, |a| a.len()) + } + + fn __len__(&self) -> PyResult { + let shape = dispatch!(&self.0, |a| a.shape()); + shape + .first() + .copied() + .ok_or_else(|| PyTypeError::new_err("len() of a rank-zero NdArray")) + } + + fn __repr__(&self) -> String { + let shape = dispatch!(&self.0, |a| a.shape().to_vec()); + format!("NdArray(shape={:?}, dtype={})", shape, self.dtype()) + } + + /// NumPy-style positional indexing. A full integer index returns a + /// scalar. Any slice or omitted trailing axis returns another NdArray + /// backed by a zero-copy view. + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { + let shape = dispatch!(&self.0, |a| a.shape().to_vec()); + let keys = axis_keys(key, &shape)?; + if keys.iter().all(|key| matches!(key, AxisKey::Index(_))) { + let indices: Vec = keys + .iter() + .map(|key| match key { + AxisKey::Index(index) => *index, + AxisKey::Range(_) => unreachable!(), + }) + .collect(); + let value = dispatch!(&self.0, |a| a.get(&indices) as f64); + return value.into_py_any(py); + } + + let selection = selectors(&keys); + let view = match &self.0 { + PyNdArrayInner::F32(a) => PyNdArray::from(a.slice(&selection)), + PyNdArrayInner::F64(a) => PyNdArray::from(a.slice(&selection)), + PyNdArrayInner::F32View(v) => PyNdArray::from(v.slice(&selection)), + PyNdArrayInner::F64View(v) => PyNdArray::from(v.slice(&selection)), + }; + Ok(Py::new(py, view)?.into_any()) + } + + /// Return a zero-copy view with axes permuted. With no argument the + /// axis order is reversed. + #[pyo3(signature = (axes=None))] + fn transpose(&self, axes: Option>) -> PyResult { + let ndim = self.ndim(); + let axes = axes.unwrap_or_else(|| (0..ndim).rev().collect()); + if axes.len() != ndim { + return Err(PyValueError::new_err(format!( + "transpose expected {} axes, got {}", + ndim, + axes.len() + ))); + } + let valid = { + let mut sorted = axes.clone(); + sorted.sort_unstable(); + sorted == (0..ndim).collect::>() + }; + if !valid { + return Err(PyValueError::new_err(format!( + "axes must be a permutation of 0..{}", + ndim + ))); + } + Ok(match &self.0 { + PyNdArrayInner::F32(a) => PyNdArray::from(a.as_view().permute_axes(&axes)), + PyNdArrayInner::F64(a) => PyNdArray::from(a.as_view().permute_axes(&axes)), + PyNdArrayInner::F32View(v) => PyNdArray::from(v.permute_axes(&axes)), + PyNdArrayInner::F64View(v) => PyNdArray::from(v.permute_axes(&axes)), + }) + } + + /// Axis-reversed zero-copy view. + #[getter(T)] + fn t(&self) -> PyResult { + self.transpose(None) + } + + /// Materialise this array into its own compact column-major allocation. + fn copy(&self) -> Self { + match &self.0 { + PyNdArrayInner::F32(a) => PyNdArray::from(a.apply(|value| value)), + PyNdArrayInner::F64(a) => PyNdArray::from(a.apply(|value| value)), + PyNdArrayInner::F32View(v) => PyNdArray::from(v.to_ndarray()), + PyNdArrayInner::F64View(v) => PyNdArray::from(v.to_ndarray()), + } + } + + /// DLPack producer entry point. Returns a capsule the consumer owns. + /// + /// A `max_version` of major 1 or above yields the versioned capsule + /// with the read-only flag carried, and consumers that support it + /// should use it. Without it, the unversioned capsule ships for + /// consumers on the pre-1.0 protocol - that capsule has no read-only + /// flag, so shared storage is copied before export. `copy=True` + /// exports a fresh compact copy in either protocol; it is always + /// writable and is flagged `IS_COPIED` on the versioned capsule. + #[pyo3(signature = (*, stream=None, max_version=None, dl_device=None, copy=None))] + fn __dlpack__( + &self, + py: Python<'_>, + stream: Option<&Bound<'_, PyAny>>, + max_version: Option<(u32, u32)>, + dl_device: Option<(i32, i32)>, + copy: Option, + ) -> PyResult> { + export_dlpack(py, &self.0, stream, max_version, dl_device, copy) + } + + /// DLPack device query. Minarrow data lives on the CPU. + fn __dlpack_device__(&self) -> (i32, i32) { + (1, 0) + } + + /// Import from a compatible CPU f32/f64 DLPack producer, such as a NumPy + /// or PyTorch tensor, or from a raw DLPack capsule. A suitably aligned + /// buffer can be shared; otherwise the data copies into an aligned buffer. + #[staticmethod] + fn from_dlpack(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { + import_dlpack(py, obj).map(PyNdArray) + } + + /// Hand to NumPy as an `ndarray` via the capsule protocol. + fn to_numpy(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let numpy = py.import("numpy")?; + Ok(numpy.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to PyTorch as a `torch.Tensor` via the capsule protocol. + fn to_pytorch(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let torch = py.import("torch")?; + Ok(torch.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to JAX as a `jax.Array` via the capsule protocol. + fn to_jax(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let jax_numpy = py.import("jax.numpy")?; + Ok(jax_numpy.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to TensorFlow as a `tf.Tensor`. TensorFlow's DLPack entry + /// takes the capsule itself rather than the producer object. + fn to_tensorflow(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let capsule = slf.borrow().__dlpack__(py, None, None, None, None)?; + let dlpack = py.import("tensorflow.experimental.dlpack")?; + Ok(dlpack.call_method1("from_dlpack", (capsule,))?.unbind()) + } + + /// Hand to CuPy via the capsule protocol. CuPy determines whether the + /// host tensor is supported and whether a device transfer is required. + fn to_cupy(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let cupy = py.import("cupy")?; + Ok(cupy.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Convert a 2D f64 tensor to a `Table` with generated column names. + fn to_table(&self) -> PyResult { + match &self.0 { + PyNdArrayInner::F64(a) => (**a) + .clone() + .to_table(None) + .map(|t| PyTable(PyTableInner::from(t))) + .map_err(|e| PyValueError::new_err(e.to_string())), + PyNdArrayInner::F64View(v) => v + .to_ndarray() + .to_table(None) + .map(|t| PyTable(PyTableInner::from(t))) + .map_err(|e| PyValueError::new_err(e.to_string())), + PyNdArrayInner::F32(_) | PyNdArrayInner::F32View(_) => Err(PyValueError::new_err( + "to_table supports float64 tensors, rebuild the tensor as float64 first", + )), + } + } +} + +/// Build a PyNdArray from a table of numeric columns. +pub fn ndarray_from_table(table: &Table) -> PyResult { + NdArray::::try_from(table) + .map(PyNdArray::from) + .map_err(|e| PyValueError::new_err(e.to_string())) +} diff --git a/minarrow-py/src/pyliquid.rs b/minarrow-py/src/pyliquid.rs index 6528477..204656f 100644 --- a/minarrow-py/src/pyliquid.rs +++ b/minarrow-py/src/pyliquid.rs @@ -50,6 +50,10 @@ //! stream whose schema is a struct, such as a RecordBatch or DataFrame. //! - `minarrow.ChunkedTable` becomes [`Value::SuperTable`]. //! - `minarrow.ChunkedArray` becomes [`Value::SuperArray`]. +//! - `minarrow.NdArray` and compatible DLPack producers without an Arrow interface, +//! such as a PyTorch tensor or NumPy `ndarray`, become `Value::NdArray` +//! under the `ndarray` feature. `Value` pins its tensors to f64, so f32 +//! tensors upcast on the way in. //! - Lists become [`Value::VecValue`] and are classified recursively. //! - Tuples of two to six elements become the corresponding fixed-arity tuple //! variant and are classified recursively. @@ -57,8 +61,6 @@ //! An empty tuple becomes `Scalar::Null`, and a one-element tuple is unwrapped. //! Tuples longer than six elements become [`Value::VecValue`]. //! -//! NumPy `ndarray` values are not currently supported directly. -//! //! ## Linking //! //! Enable the `embed` feature when embedding Python in a Rust executable. This @@ -71,7 +73,11 @@ use std::sync::Arc; use std::sync::Once; use minarrow::ffi::arrow_c_ffi::{ArrowArrayStream, ArrowSchema}; +#[cfg(feature = "ndarray")] +use minarrow::{NdArray, Vec64}; use minarrow::{Array, Scalar, Table, Value}; +#[cfg(feature = "ndarray")] +use minarrow_pyo3::ffi::dlpack::import_dlpack; use minarrow_pyo3::ffi::to_rust; use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; @@ -80,6 +86,8 @@ use pyo3::types::{PyBool, PyList, PyTuple}; use crate::array::{PyArray, PyArrayInner}; use crate::chunked_array::PyChunkedArray; use crate::chunked_table::PyChunkedTable; +#[cfg(feature = "ndarray")] +use crate::ndarray::{PyNdArray, PyNdArrayInner}; use crate::table::{PyTable, PyTableInner}; /// Resident embedded-Python runtime. @@ -102,7 +110,7 @@ impl PyLiquid { INIT.call_once(|| { crate::append_minarrow_to_inittab(); - pyo3::prepare_freethreaded_python(); + pyo3::Python::initialize(); }); PyLiquid { @@ -119,7 +127,7 @@ impl PyLiquid { I: IntoIterator, S: AsRef, { - Python::with_gil(|py| -> PyResult<()> { + Python::attach(|py| -> PyResult<()> { for name in names { let module = py.import(name.as_ref())?; self.modules.push(module.unbind()); @@ -141,7 +149,7 @@ impl PyLiquid { D: PyInput, F: for<'py> FnOnce(Python<'py>, Bound<'py, PyAny>) -> PyResult>, { - Python::with_gil(|py| { + Python::attach(|py| { let obj = data.to_python(py)?; let out = f(py, obj)?; classify(py, &out) @@ -169,6 +177,14 @@ impl PyInput for Array { } } +#[cfg(feature = "ndarray")] +impl PyInput for NdArray { + fn to_python<'py>(&self, py: Python<'py>) -> PyResult> { + let obj = PyNdArray(PyNdArrayInner::from(self.clone())); + Ok(Bound::new(py, obj)?.into_any()) + } +} + /// Converts a supported Python object to a Minarrow [`Value`]. /// /// Lists and tuples are processed recursively. Arrow-compatible objects are @@ -197,21 +213,21 @@ fn classify(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { // Check Minarrow classes before testing the generic Arrow interfaces so the // exact container variant is retained. - if obj.downcast::().is_ok() { + if obj.cast::().is_ok() { return Ok(Value::Array(Arc::new(to_rust::array_to_rust(obj)?.array))); } - if obj.downcast::().is_ok() { + if obj.cast::().is_ok() { return Ok(Value::SuperTable(Arc::new(to_rust::table_to_rust(obj)?))); } - if obj.downcast::().is_ok() { + if obj.cast::().is_ok() { return Ok(Value::SuperArray(Arc::new(to_rust::chunked_array_to_rust( obj, )?))); } - if obj.downcast::().is_ok() { + if obj.cast::().is_ok() { return Ok(Value::Table(Arc::new(to_rust::record_batch_to_rust(obj)?))); } @@ -231,7 +247,31 @@ fn classify(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { return Ok(Value::Array(Arc::new(to_rust::array_to_rust(obj)?.array))); } - if let Ok(list) = obj.downcast::() { + // Tensors from compatible DLPack producers, including `minarrow.NdArray`, + // PyTorch tensors, and NumPy arrays. `Value` pins its tensors to f64, so + // f32 data upcasts. + #[cfg(feature = "ndarray")] + if obj.cast::().is_ok() || obj.hasattr("__dlpack__")? { + return match import_dlpack(py, obj)? { + PyNdArrayInner::F64(a) => Ok(Value::NdArray(a)), + PyNdArrayInner::F64View(v) => { + Ok(Value::NdArray(Arc::new(v.to_ndarray()))) + } + PyNdArrayInner::F32(a) => { + let flat: Vec64 = (&*a).into_iter().map(|v| v as f64).collect(); + Ok(Value::NdArray(Arc::new(NdArray::from_slice(&flat, a.shape())))) + } + PyNdArrayInner::F32View(v) => { + let flat: Vec64 = (&*v).into_iter().map(|value| value as f64).collect(); + Ok(Value::NdArray(Arc::new(NdArray::from_slice( + &flat, + v.shape(), + )))) + } + }; + } + + if let Ok(list) = obj.cast::() { let mut items = Vec::with_capacity(list.len()); for item in list.iter() { @@ -241,7 +281,7 @@ fn classify(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { return Ok(Value::VecValue(Arc::new(items))); } - if let Ok(tuple) = obj.downcast::() { + if let Ok(tuple) = obj.cast::() { return classify_tuple(py, tuple); } diff --git a/minarrow-py/src/table.rs b/minarrow-py/src/table.rs index e9734b6..9c01e98 100644 --- a/minarrow-py/src/table.rs +++ b/minarrow-py/src/table.rs @@ -35,6 +35,8 @@ use pyo3::types::{PyDict, PySlice, PyTuple}; use crate::array::PyArray; use crate::convert::{build_array, build_array_typed, resolve_index}; use crate::dtype::{dtype_from_arrow, DType}; +#[cfg(feature = "ndarray")] +use crate::ndarray::{ndarray_from_table, PyNdArray}; use crate::field::PySchema; /// The natural minarrow form behind a Python `Table`. Carries the whole data @@ -71,7 +73,7 @@ impl PyTableInner { let view = self.as_view(); let n = view.n_rows(); - if let Ok(slice) = key.downcast::() { + if let Ok(slice) = key.cast::() { let ind = slice.indices(n as isize)?; if ind.step == 1 { return Ok(view.r((ind.start as usize)..(ind.stop as usize))); @@ -221,10 +223,10 @@ impl PyTableInner { pub fn get_item( &self, key: &Bound<'_, PyAny>, - wrap_array: &dyn Fn(FieldArray) -> PyResult, - wrap_table: &dyn Fn(TableV) -> PyResult, - ) -> PyResult { - if let Ok(tuple) = key.downcast::() { + wrap_array: &dyn Fn(FieldArray) -> PyResult>, + wrap_table: &dyn Fn(TableV) -> PyResult>, + ) -> PyResult> { + if let Ok(tuple) = key.cast::() { if tuple.len() == 2 { let rows = self.row_view(&tuple.get_item(0)?)?; return select_columns(&rows, &tuple.get_item(1)?, wrap_array, wrap_table); @@ -297,8 +299,8 @@ pub fn build_table_typed(data: &Bound<'_, PyDict>, schema: &Schema) -> PyResult< fn column_by_name( view: &TableV, name: &str, - wrap_array: &dyn Fn(FieldArray) -> PyResult, -) -> PyResult { + wrap_array: &dyn Fn(FieldArray) -> PyResult>, +) -> PyResult> { match view.get(name) { Some(field) => wrap_array(field), None => Err(PyKeyError::new_err(format!("column '{}' not found", name))), @@ -310,16 +312,16 @@ fn column_by_name( fn select_columns( view: &TableV, column_selection: &Bound<'_, PyAny>, - wrap_array: &dyn Fn(FieldArray) -> PyResult, - wrap_table: &dyn Fn(TableV) -> PyResult, -) -> PyResult { + wrap_array: &dyn Fn(FieldArray) -> PyResult>, + wrap_table: &dyn Fn(TableV) -> PyResult>, +) -> PyResult> { if let Ok(name) = column_selection.extract::() { return column_by_name(view, &name, wrap_array); } if let Ok(names) = column_selection.extract::>() { return wrap_table(view.c(names)); } - if let Ok(slice) = column_selection.downcast::() { + if let Ok(slice) = column_selection.cast::() { let ind = slice.indices(view.n_cols() as isize)?; let mut idxs = Vec::with_capacity(ind.slicelength); for k in 0..ind.slicelength { @@ -504,6 +506,17 @@ impl PyTable { Self::from_arrow(obj) } + /// Convert numeric columns to a 2D float64 `NdArray`, one column per + /// axis-1 entry. From there `to_numpy`, `to_pytorch`, and the other + /// DLPack bridges hand the data to each framework. + #[cfg(feature = "ndarray")] + fn to_ndarray(&self) -> PyResult { + match &self.0 { + PyTableInner::Owned(table) => ndarray_from_table(table), + PyTableInner::View(view) => ndarray_from_table(&view.to_table()), + } + } + /// Convert to a cuDF `DataFrame` through the Arrow PyCapsule interface. Runs /// on GPU and requires the `cudf` package. #[cfg(feature = "arrow_interop")] @@ -639,7 +652,7 @@ impl PyTable { self.0.repr() } - fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult { + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { self.0.get_item( key, &|field| Ok(Py::new(py, PyArray::from(field))?.into_any()), @@ -664,8 +677,8 @@ impl PyTable { fn __arrow_c_stream__( &self, py: Python<'_>, - requested_schema: Option, - ) -> PyResult { + requested_schema: Option>, + ) -> PyResult> { let _ = requested_schema; match &self.0 { PyTableInner::Owned(table) => to_py::table_to_stream_capsule(table, py), diff --git a/minarrow-py/src/xarray.rs b/minarrow-py/src/xarray.rs new file mode 100644 index 0000000..7ed04f8 --- /dev/null +++ b/minarrow-py/src/xarray.rs @@ -0,0 +1,418 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! `XArray` - labelled NdArray data for Python. + +use std::sync::Arc; + +use minarrow::structs::xarray::{Axis, XArray}; +use minarrow::{ArrayV, NdArray, NdArrayV}; +use minarrow_pyo3::ffi::dlpack::{PyNdArrayInner, export_dlpack}; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use crate::array::{PyArray, PyArrayInner}; +use crate::convert::py_to_scalar; +use crate::ndarray::{AxisKey, PyNdArray, axis_keys, selectors}; + +pub enum PyXArrayInner { + F32(Arc>), + F64(Arc>), +} + +#[pyclass(name = "XArray", module = "minarrow")] +pub struct PyXArray(pub PyXArrayInner); + +impl From> for PyXArray { + fn from(array: XArray) -> Self { + Self(PyXArrayInner::F32(Arc::new(array))) + } +} + +impl From> for PyXArray { + fn from(array: XArray) -> Self { + Self(PyXArrayInner::F64(Arc::new(array))) + } +} + +fn xarray_f32(data: &PyNdArrayInner, dims: &[String]) -> Option> { + let axes = || dims.iter().map(Axis::named).collect(); + match data { + PyNdArrayInner::F32(a) => { + let names: Vec<&str> = dims.iter().map(String::as_str).collect(); + Some(XArray::new((**a).clone(), &names)) + } + PyNdArrayInner::F32View(v) => Some(XArray::from_view((**v).clone(), axes())), + _ => None, + } +} + +fn xarray_f64(data: &PyNdArrayInner, dims: &[String]) -> Option> { + let axes = || dims.iter().map(Axis::named).collect(); + match data { + PyNdArrayInner::F64(a) => { + let names: Vec<&str> = dims.iter().map(String::as_str).collect(); + Some(XArray::new((**a).clone(), &names)) + } + PyNdArrayInner::F64View(v) => Some(XArray::from_view((**v).clone(), axes())), + _ => None, + } +} + +fn assign_coords( + array: &mut XArray, + coords: Option<&Bound<'_, PyDict>>, +) -> PyResult<()> { + let Some(coords) = coords else { + return Ok(()); + }; + for (name, values) in coords.iter() { + let name: String = name.extract()?; + let values: PyRef<'_, PyArray> = values.extract()?; + let dim = array.try_dim(&name).ok_or_else(|| { + PyValueError::new_err(format!("coordinate name '{name}' is not a dimension")) + })?; + let expected = array.shape()[dim]; + let actual = values.0.len(); + if actual != expected { + return Err(PyValueError::new_err(format!( + "coordinate '{name}' has length {actual}, expected {expected}" + ))); + } + array.assign_coords(&name, ArrayV::from(&values.0).to_array()); + } + Ok(()) +} + +#[pymethods] +impl PyXArray { + /// Label an NdArray with dimension names and optional coordinate arrays. + #[new] + #[pyo3(signature = (data, dims, coords=None))] + fn new( + data: PyRef<'_, PyNdArray>, + dims: Vec, + coords: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + let ndim = match &data.0 { + PyNdArrayInner::F32(a) => a.ndim(), + PyNdArrayInner::F64(a) => a.ndim(), + PyNdArrayInner::F32View(v) => v.ndim(), + PyNdArrayInner::F64View(v) => v.ndim(), + }; + if dims.len() != ndim { + return Err(PyValueError::new_err(format!( + "XArray needs {ndim} dimension names, got {}", + dims.len() + ))); + } + let mut unique = dims.clone(); + unique.sort(); + unique.dedup(); + if unique.len() != dims.len() { + return Err(PyValueError::new_err( + "XArray dimension names must be unique", + )); + } + if let Some(mut array) = xarray_f32(&data.0, &dims) { + assign_coords(&mut array, coords)?; + return Ok(array.into()); + } + let mut array = xarray_f64(&data.0, &dims).ok_or_else(|| { + PyValueError::new_err("XArray data must have dtype float32 or float64") + })?; + assign_coords(&mut array, coords)?; + Ok(array.into()) + } + + #[getter] + fn shape<'py>(&self, py: Python<'py>) -> PyResult> { + match &self.0 { + PyXArrayInner::F32(a) => PyTuple::new(py, a.shape()), + PyXArrayInner::F64(a) => PyTuple::new(py, a.shape()), + } + } + + #[getter] + fn ndim(&self) -> usize { + match &self.0 { + PyXArrayInner::F32(a) => a.ndim(), + PyXArrayInner::F64(a) => a.ndim(), + } + } + + #[getter] + fn size(&self) -> usize { + match &self.0 { + PyXArrayInner::F32(a) => a.len(), + PyXArrayInner::F64(a) => a.len(), + } + } + + #[getter] + fn dtype(&self) -> &'static str { + match &self.0 { + PyXArrayInner::F32(_) => "float32", + PyXArrayInner::F64(_) => "float64", + } + } + + #[getter] + fn dims(&self) -> Vec { + match &self.0 { + PyXArrayInner::F32(a) => a.dim_names().into_iter().map(str::to_string).collect(), + PyXArrayInner::F64(a) => a.dim_names().into_iter().map(str::to_string).collect(), + } + } + + #[getter] + fn coords<'py>(&self, py: Python<'py>) -> PyResult> { + let result = PyDict::new(py); + match &self.0 { + PyXArrayInner::F32(a) => { + for axis in a.axes() { + if let Some(coords) = &axis.coords { + result.set_item( + &axis.name, + Py::new(py, PyArray(PyArrayInner::from(coords.clone())))?, + )?; + } + } + } + PyXArrayInner::F64(a) => { + for axis in a.axes() { + if let Some(coords) = &axis.coords { + result.set_item( + &axis.name, + Py::new(py, PyArray(PyArrayInner::from(coords.clone())))?, + )?; + } + } + } + } + Ok(result) + } + + /// The underlying data as an NdArray. If this XArray is itself a + /// selection, the returned NdArray remains a zero-copy view. + #[getter] + fn data(&self) -> PyNdArray { + match &self.0 { + PyXArrayInner::F32(a) if a.is_owned() => PyNdArray::from((**a).clone().into_ndarray()), + PyXArrayInner::F64(a) if a.is_owned() => PyNdArray::from((**a).clone().into_ndarray()), + PyXArrayInner::F32(a) => PyNdArray::from(a.as_view()), + PyXArrayInner::F64(a) => PyNdArray::from(a.as_view()), + } + } + + fn __len__(&self) -> PyResult { + let shape = match &self.0 { + PyXArrayInner::F32(a) => a.shape(), + PyXArrayInner::F64(a) => a.shape(), + }; + shape + .first() + .copied() + .ok_or_else(|| PyTypeError::new_err("len() of a rank-zero XArray")) + } + + fn __repr__(&self) -> String { + format!( + "XArray(shape={:?}, dims={:?}, dtype={})", + match &self.0 { + PyXArrayInner::F32(a) => a.shape(), + PyXArrayInner::F64(a) => a.shape(), + }, + self.dims(), + self.dtype() + ) + } + + /// Positional selection with NdArray indexing semantics, preserving the + /// names and coordinates of axes that remain. + fn __getitem__(&self, py: Python<'_>, key: &Bound<'_, PyAny>) -> PyResult> { + let shape = match &self.0 { + PyXArrayInner::F32(a) => a.shape(), + PyXArrayInner::F64(a) => a.shape(), + }; + let keys = axis_keys(key, &shape)?; + if keys.iter().all(|key| matches!(key, AxisKey::Index(_))) { + let indices: Vec = keys + .iter() + .map(|key| match key { + AxisKey::Index(index) => *index, + AxisKey::Range(_) => unreachable!(), + }) + .collect(); + let value = match &self.0 { + PyXArrayInner::F32(a) => a.get(&indices) as f64, + PyXArrayInner::F64(a) => a.get(&indices), + }; + return value.into_py_any(py); + } + + let selection = selectors(&keys); + let result = match &self.0 { + PyXArrayInner::F32(a) => { + let names = a.dim_names(); + let named: Vec<(&str, &dyn minarrow::traits::selection::DataSelector)> = + names.into_iter().zip(selection).collect(); + PyXArray::from(a.select(&named)) + } + PyXArrayInner::F64(a) => { + let names = a.dim_names(); + let named: Vec<(&str, &dyn minarrow::traits::selection::DataSelector)> = + names.into_iter().zip(selection).collect(); + PyXArray::from(a.select(&named)) + } + }; + Ok(Py::new(py, result)?.into_any()) + } + + /// Exact coordinate selection. + /// + /// Behaviour: + /// - Each keyword is a dimension name + /// - Selected dimensions collapse in the result. + #[pyo3(signature = (**indexers))] + fn sel(&self, indexers: Option<&Bound<'_, PyDict>>) -> PyResult { + let indexers = indexers.ok_or_else(|| PyValueError::new_err("sel needs an indexer"))?; + match &self.0 { + PyXArrayInner::F32(a) => { + let mut result = (**a).clone(); + for (name, value) in indexers.iter() { + result = result + .try_at(&name.extract::()?, py_to_scalar(&value)?) + .map_err(|e| PyIndexError::new_err(e.to_string()))?; + } + Ok(result.into()) + } + PyXArrayInner::F64(a) => { + let mut result = (**a).clone(); + for (name, value) in indexers.iter() { + result = result + .try_at(&name.extract::()?, py_to_scalar(&value)?) + .map_err(|e| PyIndexError::new_err(e.to_string()))?; + } + Ok(result.into()) + } + } + } + + /// Inclusive coordinate range selection on one dimension. + fn between( + &self, + dim: &str, + low: &Bound<'_, PyAny>, + high: &Bound<'_, PyAny>, + ) -> PyResult { + let low = py_to_scalar(low)?; + let high = py_to_scalar(high)?; + match &self.0 { + PyXArrayInner::F32(a) => a + .try_between(dim, low, high) + .map(PyXArray::from) + .map_err(|e| PyIndexError::new_err(e.to_string())), + PyXArrayInner::F64(a) => a + .try_between(dim, low, high) + .map(PyXArray::from) + .map_err(|e| PyIndexError::new_err(e.to_string())), + } + } + + /// Closest coordinate selection on one numeric or datetime dimension. + fn nearest(&self, dim: &str, value: &Bound<'_, PyAny>) -> PyResult { + let value = py_to_scalar(value)?; + match &self.0 { + PyXArrayInner::F32(a) => a + .try_nearest(dim, value) + .map(PyXArray::from) + .map_err(|e| PyIndexError::new_err(e.to_string())), + PyXArrayInner::F64(a) => a + .try_nearest(dim, value) + .map(PyXArray::from) + .map_err(|e| PyIndexError::new_err(e.to_string())), + } + } + + #[pyo3(signature = (*, stream=None, max_version=None, dl_device=None, copy=None))] + fn __dlpack__( + &self, + py: Python<'_>, + stream: Option<&Bound<'_, PyAny>>, + max_version: Option<(u32, u32)>, + dl_device: Option<(i32, i32)>, + copy: Option, + ) -> PyResult> { + match &self.0 { + PyXArrayInner::F32(a) => export_dlpack( + py, + &PyNdArrayInner::from(a.as_view()), + stream, + max_version, + dl_device, + copy, + ), + PyXArrayInner::F64(a) => export_dlpack( + py, + &PyNdArrayInner::from(a.as_view()), + stream, + max_version, + dl_device, + copy, + ), + } + } + + fn __dlpack_device__(&self) -> (i32, i32) { + (1, 0) + } + + fn to_numpy(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let numpy = py.import("numpy")?; + Ok(numpy.call_method1("from_dlpack", (slf,))?.unbind()) + } +} + +impl From> for PyXArray { + fn from(array: NdArray) -> Self { + XArray::from_ndarray(array).into() + } +} + +impl From> for PyXArray { + fn from(array: NdArray) -> Self { + XArray::from_ndarray(array).into() + } +} + +impl From> for PyXArray { + fn from(view: NdArrayV) -> Self { + let axes = (0..view.ndim()) + .map(|i| Axis::named(format!("dim_{i}"))) + .collect(); + XArray::from_view(view, axes).into() + } +} + +impl From> for PyXArray { + fn from(view: NdArrayV) -> Self { + let axes = (0..view.ndim()) + .map(|i| Axis::named(format!("dim_{i}"))) + .collect(); + XArray::from_view(view, axes).into() + } +} diff --git a/minarrow-py/tests/test_ndarray.py b/minarrow-py/tests/test_ndarray.py new file mode 100644 index 0000000..7eb284a --- /dev/null +++ b/minarrow-py/tests/test_ndarray.py @@ -0,0 +1,426 @@ +"""Tests for the native minarrow-py NdArray surface and its DLPack bridges. + +Run after `maturin develop`: + python -m pytest tests/test_ndarray.py + +The framework bridge tests skip when the target library is not installed. +""" + +import pytest + +import minarrow as mp + + +# --- Construction and inspection ------------------------------------------- + + +def test_construct_and_inspect(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + assert a.shape == (3, 2) + assert a.strides == (1, 3) + assert a.ndim == 2 + assert a.size == 6 + assert a.dtype == "float64" + assert len(a) == 3 + assert repr(a) == "NdArray(shape=[3, 2], dtype=float64)" + + +def test_construct_1d_default_shape(): + a = mp.NdArray([1.0, 2.0, 3.0]) + assert a.shape == (3,) + assert a[1] == 2.0 + + +def test_construct_f32(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2], dtype="float32") + assert a.dtype == "float32" + assert a[1, 1] == 4.0 + + +def test_construct_shape_mismatch(): + with pytest.raises(ValueError): + mp.NdArray([1.0, 2.0, 3.0], shape=[2, 2]) + + +def test_construct_bad_dtype(): + with pytest.raises(ValueError): + mp.NdArray([1.0], dtype="int64") + + +def test_getitem_column_major(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + assert a[0, 0] == 1.0 + assert a[2, 0] == 3.0 + assert a[0, 1] == 4.0 + assert a[2, 1] == 6.0 + + +def test_getitem_out_of_bounds(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + with pytest.raises(IndexError): + a[2, 0] + row = a[0] + assert isinstance(row, mp.NdArray) + assert row.shape == (2,) + assert row[1] == 3.0 + + +def test_getitem_returns_zero_copy_views(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + v = a[1:, :] + assert isinstance(v, mp.NdArray) + assert v.is_view is True + assert v.shape == (2, 2) + assert v.strides == (1, 3) + assert v[0, 0] == 2.0 + assert v[-1, -1] == 6.0 + + +def test_transpose_is_same_python_type(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + t = a.T + assert isinstance(t, mp.NdArray) + assert t.is_view is True + assert t.shape == (2, 3) + assert t.strides == (3, 1) + assert t[1, 2] == 6.0 + + +# --- DLPack protocol -------------------------------------------------------- + + +def test_dlpack_device(): + a = mp.NdArray([1.0, 2.0]) + assert a.__dlpack_device__() == (1, 0) + + +def test_dlpack_capsule_names(): + a = mp.NdArray([1.0, 2.0, 3.0]) + legacy = a.__dlpack__() + versioned = a.__dlpack__(max_version=(1, 1)) + # repr quotes the exact capsule name, so match it in full to keep + # "dltensor" from passing on a "dltensor_versioned" capsule. + assert '"dltensor"' in repr(legacy) + assert '"dltensor_versioned"' in repr(versioned) + + +def test_rank_zero_scalar_and_dlpack(): + a = mp.NdArray([5.0], shape=[]) + assert a.shape == () + assert a.strides == () + assert a.ndim == 0 + assert a.size == 1 + assert a[()] == 5.0 + assert a.T.shape == () + with pytest.raises(TypeError): + len(a) + + b = mp.NdArray.from_dlpack(a) + assert b.shape == () + assert b[()] == 5.0 + + np = pytest.importorskip("numpy") + n = np.from_dlpack(a) + assert n.shape == () + assert n.item() == 5.0 + + +def test_dlpack_rejects_foreign_device(): + a = mp.NdArray([1.0, 2.0]) + with pytest.raises(ValueError): + a.__dlpack__(dl_device=(2, 0)) + + +def test_from_dlpack_roundtrip_self(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + b = mp.NdArray.from_dlpack(a) + assert b.shape == (3, 2) + assert b[2, 1] == 6.0 + assert b.dtype == "float64" + + +def test_from_dlpack_rejects_non_producer(): + with pytest.raises((ValueError, AttributeError)): + mp.NdArray.from_dlpack(object()) + + +# --- NumPy bridge ----------------------------------------------------------- + + +def test_numpy_roundtrip(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + v = np.from_dlpack(a) + assert v.shape == (3, 2) + assert v[0, 0] == 1.0 + assert v[2, 1] == 6.0 + + back = mp.NdArray.from_dlpack(np.arange(6, dtype="float64").reshape(2, 3)) + assert back.shape == (2, 3) + assert back[0, 0] == 0.0 + assert back[1, 2] == 5.0 + + +def test_to_numpy_method(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + v = a.to_numpy() + assert isinstance(v, np.ndarray) + assert v.dtype == np.float64 + assert v[1, 1] == 4.0 + + +def test_numpy_f32_roundtrip(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0], dtype="float32") + v = a.to_numpy() + assert v.dtype == np.float32 + back = mp.NdArray.from_dlpack(np.asarray([1.5, 2.5], dtype="float32")) + assert back.dtype == "float32" + assert back[1] == 2.5 + + +def test_dlpack_copy_is_writable(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0]) + v = np.from_dlpack(a, copy=True) + assert v.flags.writeable is True + v[0] = 99.0 + # The copy is independent of the source tensor. + assert a[0] == 1.0 + + +def test_dlpack_shared_view_is_read_only(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0]) + v = np.from_dlpack(a) + # The versioned capsule shares the buffer, so it arrives read-only. + assert v.flags.writeable is False + + +def test_dlpack_ndarray_slice_preserves_shape_and_values(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + v = np.from_dlpack(a[1:, :]) + assert v.shape == (2, 2) + assert v[0, 0] == 2.0 + assert v[1, 1] == 6.0 + + +def test_dlpack_export_shares_one_address(): + np = pytest.importorskip("numpy") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + p1 = np.from_dlpack(a).__array_interface__["data"][0] + p2 = np.from_dlpack(a).__array_interface__["data"][0] + assert p1 == p2 + # An explicit copy detaches to a fresh allocation. + p3 = np.from_dlpack(a, copy=True).__array_interface__["data"][0] + assert p3 != p1 + + +def test_dlpack_aligned_import_wraps_source_memory(): + np = pytest.importorskip("numpy") + # Carve a 64-byte aligned window out of an over-allocated buffer so + # the import takes the zero-copy branch deterministically. + raw = np.zeros(16 + 8, dtype=np.float64) + base = raw.__array_interface__["data"][0] + off = (-base) % 64 // 8 + src = raw[off:off + 16].reshape(4, 4) + src[:] = np.arange(16, dtype=np.float64).reshape(4, 4) + p_src = src.__array_interface__["data"][0] + assert p_src % 64 == 0 + + b = mp.NdArray.from_dlpack(src) + p_rt = np.from_dlpack(b).__array_interface__["data"][0] + assert p_rt == p_src + + # A misaligned window copies into a fresh 64-byte aligned buffer. + mis = raw[off + 1:off + 9] + c = mp.NdArray.from_dlpack(mis) + p_c = np.from_dlpack(c).__array_interface__["data"][0] + assert p_c != mis.__array_interface__["data"][0] + assert p_c % 64 == 0 + + +# --- Framework bridges (skip when the library is absent) -------------------- + + +def test_to_pytorch(): + torch = pytest.importorskip("torch") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + t = a.to_pytorch() + assert isinstance(t, torch.Tensor) + assert t.shape == (2, 2) + assert t[1, 1].item() == 4.0 + + +def test_to_jax(): + pytest.importorskip("jax") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + t = a.to_jax() + assert t.shape == (2, 2) + assert float(t[1, 1]) == 4.0 + + +def test_to_tensorflow(): + pytest.importorskip("tensorflow") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + t = a.to_tensorflow() + assert tuple(t.shape) == (2, 2) + + +def test_to_cupy(): + cupy = pytest.importorskip("cupy") + a = mp.NdArray([1.0, 2.0, 3.0, 4.0], shape=[2, 2]) + t = a.to_cupy() + assert t.shape == (2, 2) + + +# --- Table interop ---------------------------------------------------------- + + +def test_table_to_ndarray(): + t = mp.Table({"x": [1.0, 2.0, 3.0], "y": [4.0, 5.0, 6.0]}) + a = t.to_ndarray() + assert a.shape == (3, 2) + assert a[0, 0] == 1.0 + assert a[2, 1] == 6.0 + + +def test_ndarray_to_table(): + a = mp.NdArray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2]) + t = a.to_table() + assert t.n_rows == 3 + assert t.n_cols == 2 + + +# --- ChunkedNdArray --------------------------------------------------------- + + +def make_chunked_ndarray(): + a = mp.NdArray([1.0, 2.0, 10.0, 20.0], shape=[2, 2]) + b = mp.NdArray([3.0, 30.0], shape=[1, 2]) + return mp.ChunkedNdArray([a, b], name="readings") + + +def test_chunked_ndarray_construction_and_access(): + a = make_chunked_ndarray() + assert a.shape == (3, 2) + assert a.ndim == 2 + assert a.size == 6 + assert a.dtype == "float64" + assert a.name == "readings" + assert a.n_chunks == 2 + assert a.is_view is False + assert isinstance(a.chunk(0), mp.NdArray) + assert a.chunk(2) is None + assert len(a) == 3 + assert a[2, 1] == 30.0 + row = a[1] + assert isinstance(row, mp.NdArray) + assert row.shape == (2,) + assert row[1] == 20.0 + + +def test_chunked_ndarray_rejects_rank_zero_pieces(): + with pytest.raises(ValueError, match="axis 0"): + mp.ChunkedNdArray([mp.NdArray([5.0], shape=[])]) + + +def test_chunked_ndarray_slice_stays_chunked(): + a = make_chunked_ndarray() + v = a[1:] + assert isinstance(v, mp.ChunkedNdArray) + assert v.shape == (2, 2) + assert v.n_chunks == 2 + assert v.is_view is True + assert v[0, 0] == 2.0 + assert v[1, 1] == 30.0 + assert all(isinstance(chunk, mp.NdArray) for chunk in v.chunks) + + +def test_chunked_ndarray_materialises_in_logical_column_major_order(): + a = make_chunked_ndarray().to_ndarray() + assert isinstance(a, mp.NdArray) + assert a.shape == (3, 2) + assert [a[i, 0] for i in range(3)] == [1.0, 2.0, 3.0] + assert [a[i, 1] for i in range(3)] == [10.0, 20.0, 30.0] + + +def test_chunked_ndarray_dlpack_is_per_chunk(): + a = make_chunked_ndarray() + assert not hasattr(a, "__dlpack__") + + chunks = a.chunks + assert len(chunks) == 2 + assert all(hasattr(chunk, "__dlpack__") for chunk in chunks) + + arrays = a.to_numpy() + assert [array.shape for array in arrays] == [(2, 2), (1, 2)] + assert arrays[0].tolist() == [[1.0, 10.0], [2.0, 20.0]] + assert arrays[1].tolist() == [[3.0, 30.0]] + + pointers = [array.__array_interface__["data"][0] for array in arrays] + assert pointers[0] != pointers[1] + + +# --- XArray ---------------------------------------------------------------- + + +def make_xarray(): + data = mp.NdArray([1.0, 2.0, 3.0, 10.0, 20.0, 30.0], shape=[3, 2]) + return mp.XArray( + data, + dims=["time", "feature"], + coords={ + "time": mp.Array([10, 20, 30]), + "feature": mp.Array(["x", "y"]), + }, + ) + + +def test_xarray_construction_and_positional_selection(): + a = make_xarray() + assert a.shape == (3, 2) + assert a.dims == ["time", "feature"] + assert set(a.coords) == {"time", "feature"} + assert a.data.is_view is False + assert a[2, 1] == 30.0 + row = a[1] + assert isinstance(row, mp.XArray) + assert row.shape == (2,) + assert row.dims == ["feature"] + assert row.data.is_view is True + assert row.data[1] == 20.0 + + +def test_rank_zero_xarray(): + a = mp.XArray(mp.NdArray([5.0], shape=[]), dims=[]) + assert a.shape == () + assert a.dims == [] + assert a.coords == {} + assert a[()] == 5.0 + with pytest.raises(TypeError): + len(a) + + +def test_xarray_coordinate_selection(): + a = make_xarray() + exact = a.sel(time=20) + assert exact.dims == ["feature"] + assert exact.data[0] == 2.0 + assert exact.data[1] == 20.0 + + window = a.between("time", 15, 30) + assert window.shape == (2, 2) + assert window.data[0, 0] == 2.0 + assert window.data[1, 1] == 30.0 + + nearest = a.nearest("time", 19) + assert nearest.dims == ["feature"] + assert nearest.data[1] == 20.0 + + +def test_xarray_dlpack_exports_data_only(): + a = make_xarray() + assert '"dltensor_versioned"' in repr(a.__dlpack__(max_version=(1, 1))) diff --git a/pyo3/Cargo.lock b/pyo3/Cargo.lock index 77a177e..424d771 100644 --- a/pyo3/Cargo.lock +++ b/pyo3/Cargo.lock @@ -8,27 +8,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "cfg-if" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" - [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "libc" version = "0.2.178" @@ -41,15 +26,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "minarrow" version = "0.15.0" @@ -102,37 +78,32 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-build-config" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -140,9 +111,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -152,13 +123,12 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.23.5" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] @@ -172,12 +142,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - [[package]] name = "syn" version = "2.0.111" @@ -191,9 +155,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "thiserror" @@ -221,12 +185,6 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "vec64" version = "0.4.7" diff --git a/pyo3/Cargo.toml b/pyo3/Cargo.toml index 9e116aa..f3b451c 100644 --- a/pyo3/Cargo.toml +++ b/pyo3/Cargo.toml @@ -25,7 +25,7 @@ crate-type = ["cdylib", "rlib"] [dependencies] minarrow = { version = "0.15.0", path = "..", features = ["large_string"] } -pyo3 = { version = "0.23" } +pyo3 = { version = "0.29" } thiserror = "2" [features] @@ -35,9 +35,11 @@ datetime = ["minarrow/datetime"] extended_numeric_types = ["minarrow/extended_numeric_types"] extended_categorical = ["minarrow/extended_categorical"] table_metadata = ["minarrow/table_metadata"] +# N-dimensional tensor bridging with DLPack capsule interchange. +ndarray = ["minarrow/ndarray", "minarrow/dlpack", "minarrow/views", "minarrow/select"] [build-dependencies] -pyo3-build-config = "0.23" +pyo3-build-config = "0.29" [[example]] name = "pycapsule_exchange" diff --git a/pyo3/examples/pycapsule_exchange.rs b/pyo3/examples/pycapsule_exchange.rs index ca87034..d8174c3 100644 --- a/pyo3/examples/pycapsule_exchange.rs +++ b/pyo3/examples/pycapsule_exchange.rs @@ -75,9 +75,9 @@ use pyo3::types::IntoPyDict; use std::sync::Arc; fn main() -> PyResult<()> { - pyo3::prepare_freethreaded_python(); + Python::initialize(); - Python::with_gil(|py| { + Python::attach(|py| { println!("=== Arrow PyCapsule Interface Examples ===\n"); example_1_export_array(py)?; diff --git a/pyo3/src/ffi/dlpack.rs b/pyo3/src/ffi/dlpack.rs new file mode 100644 index 0000000..86703c7 --- /dev/null +++ b/pyo3/src/ffi/dlpack.rs @@ -0,0 +1,322 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! # DLPack Capsule Interchange +//! +//! The Python-boundary half of minarrow's DLPack support: capsule +//! construction and consumption for the `__dlpack__` protocol, shared by +//! this crate's `PyNdArray` and by minarrow-py's `NdArray`, mirroring how +//! [`to_py`](crate::ffi::to_py) and [`to_rust`](crate::ffi::to_rust) host +//! the Arrow capsule glue for both crates. +//! +//! [`export_dlpack`] wraps an f32/f64 [`PyNdArrayInner`] in a legacy or +//! versioned DLPack capsule, and [`import_dlpack`] consumes any DLPack +//! producer object or raw capsule back into one. Consumed capsules rename +//! with the `used_` prefix so their destructors do not release the tensor, and +//! unconsumed capsules release the tensor through the capsule destructor. + +use std::ffi::{CStr, c_void}; +use std::sync::Arc; + +use minarrow::{NdArray, NdArrayV}; +use minarrow::ffi::dlpack::{ + DLManagedTensor, DLManagedTensorVersioned, DLPACK_FLAG_BITMASK_IS_COPIED, + DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION, export_to_dlpack, export_to_dlpack_versioned, + export_view_to_dlpack, export_view_to_dlpack_versioned, import_from_dlpack, + import_from_dlpack_versioned, +}; +use minarrow::traits::type_unions::Float; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +// Capsule names from the DLPack Python protocol. A consumed capsule +// renames with the `used_` prefix so its destructor does not free a +// tensor the consumer now owns. +pub static DLTENSOR: &CStr = c"dltensor"; +pub static DLTENSOR_USED: &CStr = c"used_dltensor"; +pub static DLTENSOR_VERSIONED: &CStr = c"dltensor_versioned"; +pub static DLTENSOR_VERSIONED_USED: &CStr = c"used_dltensor_versioned"; + +/// Internal storage for a Python `NdArray`, covering both supported element +/// types and owned or windowed data. +pub enum PyNdArrayInner { + F32(Arc>), + F64(Arc>), + F32View(Arc>), + F64View(Arc>), +} + +impl From> for PyNdArrayInner { + fn from(ndarray: NdArray) -> Self { + PyNdArrayInner::F32(Arc::new(ndarray)) + } +} + +impl From> for PyNdArrayInner { + fn from(ndarray: NdArray) -> Self { + PyNdArrayInner::F64(Arc::new(ndarray)) + } +} + +impl From> for PyNdArrayInner { + fn from(view: NdArrayV) -> Self { + PyNdArrayInner::F32View(Arc::new(view)) + } +} + +impl From> for PyNdArrayInner { + fn from(view: NdArrayV) -> Self { + PyNdArrayInner::F64View(Arc::new(view)) + } +} + +/// Destructor for an unconsumed legacy capsule. A consumed capsule is +/// renamed `used_dltensor` and the consumer owns the tensor. +unsafe extern "C" fn dltensor_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) { + if unsafe { pyo3::ffi::PyCapsule_IsValid(capsule, DLTENSOR.as_ptr()) } == 1 { + let raw = unsafe { pyo3::ffi::PyCapsule_GetPointer(capsule, DLTENSOR.as_ptr()) } + as *mut DLManagedTensor; + if !raw.is_null() { + unsafe { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + } + } + } +} + +/// Destructor for an unconsumed versioned capsule. +unsafe extern "C" fn dltensor_versioned_capsule_destructor(capsule: *mut pyo3::ffi::PyObject) { + if unsafe { pyo3::ffi::PyCapsule_IsValid(capsule, DLTENSOR_VERSIONED.as_ptr()) } == 1 { + let raw = unsafe { pyo3::ffi::PyCapsule_GetPointer(capsule, DLTENSOR_VERSIONED.as_ptr()) } + as *mut DLManagedTensorVersioned; + if !raw.is_null() { + unsafe { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + } + } + } +} + +/// The `__dlpack__` protocol body. Validates the standard keyword +/// arguments, resolves the requested ABI, and returns a capsule the +/// consumer owns. +/// +/// A `max_version` of major 1 or above yields the versioned capsule with +/// the read-only flag carried. Without it, the unversioned capsule ships +/// for consumers on the pre-1.0 protocol - that capsule has no read-only +/// flag, so shared storage is copied before export. `copy=True` exports a +/// fresh compact copy in either protocol; it is always writable and is +/// flagged `IS_COPIED` on the versioned capsule. +pub fn export_dlpack( + py: Python<'_>, + inner: &PyNdArrayInner, + stream: Option<&Bound<'_, PyAny>>, + max_version: Option<(u32, u32)>, + dl_device: Option<(i32, i32)>, + copy: Option, +) -> PyResult> { + if let Some(stream) = stream { + if !stream.is_none() { + return Err(PyValueError::new_err("stream must be None for CPU tensors")); + } + } + if let Some(device) = dl_device { + if device != (1, 0) { + return Err(PyValueError::new_err(format!( + "cannot export to device {:?}, data lives on CPU (1, 0)", + device + ))); + } + } + let versioned = matches!(max_version, Some((major, _)) if major >= 1); + let copy = copy.unwrap_or(false); + + match inner { + PyNdArrayInner::F32(a) => { + let source = if copy { a.apply(|v| v) } else { (**a).clone() }; + dlpack_capsule(py, source, versioned, copy) + } + PyNdArrayInner::F64(a) => { + let source = if copy { a.apply(|v| v) } else { (**a).clone() }; + dlpack_capsule(py, source, versioned, copy) + } + PyNdArrayInner::F32View(v) => { + if copy { + dlpack_capsule(py, v.to_ndarray(), versioned, true) + } else { + dlpack_view_capsule(py, (**v).clone(), versioned) + } + } + PyNdArrayInner::F64View(v) => { + if copy { + dlpack_capsule(py, v.to_ndarray(), versioned, true) + } else { + dlpack_view_capsule(py, (**v).clone(), versioned) + } + } + } +} + +/// Wrap an owned NdArray in a DLPack capsule of the requested ABI. +/// `copied` marks a versioned capsule with `DLPACK_FLAG_BITMASK_IS_COPIED` +/// so the consumer knows the data does not alias the exporter's memory. +pub fn dlpack_capsule( + py: Python<'_>, + source: NdArray, + versioned: bool, + copied: bool, +) -> PyResult> { + if versioned { + let raw = export_to_dlpack_versioned(source).into_raw(); + if copied { + unsafe { + (*raw).flags |= DLPACK_FLAG_BITMASK_IS_COPIED; + } + } + versioned_capsule(py, raw) + } else { + let raw = export_to_dlpack(source).into_raw(); + legacy_capsule(py, raw) + } +} + +/// Wrap a windowed NdArray in a DLPack capsule of the requested ABI. +pub fn dlpack_view_capsule( + py: Python<'_>, + view: NdArrayV, + versioned: bool, +) -> PyResult> { + if versioned { + versioned_capsule(py, export_view_to_dlpack_versioned(view).into_raw()) + } else { + legacy_capsule(py, export_view_to_dlpack(view).into_raw()) + } +} + +fn legacy_capsule(py: Python<'_>, raw: *mut DLManagedTensor) -> PyResult> { + let capsule = unsafe { + pyo3::ffi::PyCapsule_New( + raw as *mut c_void, + DLTENSOR.as_ptr(), + Some(dltensor_capsule_destructor), + ) + }; + if capsule.is_null() { + unsafe { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + } + return Err(PyErr::fetch(py)); + } + Ok(unsafe { Bound::from_owned_ptr(py, capsule) }.unbind()) +} + +fn versioned_capsule( + py: Python<'_>, + raw: *mut DLManagedTensorVersioned, +) -> PyResult> { + let capsule = unsafe { + pyo3::ffi::PyCapsule_New( + raw as *mut c_void, + DLTENSOR_VERSIONED.as_ptr(), + Some(dltensor_versioned_capsule_destructor), + ) + }; + if capsule.is_null() { + unsafe { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + } + return Err(PyErr::fetch(py)); + } + Ok(unsafe { Bound::from_owned_ptr(py, capsule) }.unbind()) +} + +/// Import from a compatible CPU f32/f64 DLPack producer, such as a NumPy or +/// PyTorch tensor, or from a raw DLPack capsule. A suitably aligned buffer can +/// be shared; otherwise the data copies into an aligned buffer. +pub fn import_dlpack(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { + let capsule = if unsafe { pyo3::ffi::PyCapsule_CheckExact(obj.as_ptr()) } == 1 { + obj.clone() + } else { + let kwargs = PyDict::new(py); + kwargs.set_item("max_version", (DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION))?; + match obj.call_method("__dlpack__", (), Some(&kwargs)) { + Ok(capsule) => capsule, + // A pre-1.0 producer rejects the max_version keyword. + Err(e) if e.is_instance_of::(py) => obj.call_method0("__dlpack__")?, + Err(e) => return Err(e), + } + }; + + let cap_ptr = capsule.as_ptr(); + unsafe { + if pyo3::ffi::PyCapsule_IsValid(cap_ptr, DLTENSOR_VERSIONED.as_ptr()) == 1 { + let raw = pyo3::ffi::PyCapsule_GetPointer(cap_ptr, DLTENSOR_VERSIONED.as_ptr()) + as *mut DLManagedTensorVersioned; + // Ownership transfers to the import, so the capsule renames + // first and its destructor stands down. + pyo3::ffi::PyCapsule_SetName(cap_ptr, DLTENSOR_VERSIONED_USED.as_ptr()); + match (*raw).dl_tensor.dtype.bits { + 32 => import_from_dlpack_versioned::(raw) + .map(PyNdArrayInner::from) + .map_err(|e| PyValueError::new_err(e.to_string())), + 64 => import_from_dlpack_versioned::(raw) + .map(PyNdArrayInner::from) + .map_err(|e| PyValueError::new_err(e.to_string())), + bits => { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + Err(PyValueError::new_err(format!( + "unsupported DLPack element width {} bits, expected 32 or 64", + bits + ))) + } + } + } else if pyo3::ffi::PyCapsule_IsValid(cap_ptr, DLTENSOR.as_ptr()) == 1 { + let raw = + pyo3::ffi::PyCapsule_GetPointer(cap_ptr, DLTENSOR.as_ptr()) as *mut DLManagedTensor; + pyo3::ffi::PyCapsule_SetName(cap_ptr, DLTENSOR_USED.as_ptr()); + match (*raw).dl_tensor.dtype.bits { + 32 => import_from_dlpack::(raw) + .map(PyNdArrayInner::from) + .map_err(|e| PyValueError::new_err(e.to_string())), + 64 => import_from_dlpack::(raw) + .map(PyNdArrayInner::from) + .map_err(|e| PyValueError::new_err(e.to_string())), + bits => { + if let Some(deleter) = (*raw).deleter { + deleter(raw); + } + Err(PyValueError::new_err(format!( + "unsupported DLPack element width {} bits, expected 32 or 64", + bits + ))) + } + } + } else { + Err(PyValueError::new_err( + "expected a DLPack capsule or an object with __dlpack__", + )) + } + } +} diff --git a/pyo3/src/ffi/mod.rs b/pyo3/src/ffi/mod.rs index 1872539..85792a4 100644 --- a/pyo3/src/ffi/mod.rs +++ b/pyo3/src/ffi/mod.rs @@ -191,6 +191,9 @@ //! //! - [`to_py`] - MinArrow to Python conversion (export) //! - [`to_rust`] - Python to MinArrow conversion (import) +//! - `dlpack` - DLPack capsule interchange for `NdArray` (`ndarray` feature) +#[cfg(feature = "ndarray")] +pub mod dlpack; pub mod to_py; pub mod to_rust; diff --git a/pyo3/src/ffi/to_py.rs b/pyo3/src/ffi/to_py.rs index 1a7b3b8..204f7df 100644 --- a/pyo3/src/ffi/to_py.rs +++ b/pyo3/src/ffi/to_py.rs @@ -482,7 +482,7 @@ pub fn array_to_capsules<'py>( array: Arc, field: &Field, py: Python<'py>, -) -> PyResult<(PyObject, PyObject)> { +) -> PyResult<(Py, Py)> { let schema = Schema::from(vec![field.clone()]); let (arr_ptr, sch_ptr) = export_to_c(array, schema); @@ -540,7 +540,7 @@ pub fn array_to_capsules<'py>( /// Exports a MinArrow Table as an ArrowArrayStream PyCapsule. /// /// The stream yields one struct array (record batch) corresponding to the table. -pub fn table_to_stream_capsule<'py>(table: &Table, py: Python<'py>) -> PyResult { +pub fn table_to_stream_capsule<'py>(table: &Table, py: Python<'py>) -> PyResult> { let fields: Vec = table.cols.iter().map(|fa| (*fa.field).clone()).collect(); let columns: Vec<(Arc, Schema)> = table .cols @@ -587,7 +587,7 @@ pub fn table_to_stream_capsule<'py>(table: &Table, py: Python<'py>) -> PyResult< pub fn super_table_to_stream_capsule<'py>( super_table: &SuperTable, py: Python<'py>, -) -> PyResult { +) -> PyResult> { if super_table.batches.is_empty() { return Err(pyo3::exceptions::PyValueError::new_err( "Cannot export empty SuperTable as stream capsule", @@ -821,7 +821,7 @@ pub fn array_view_to_capsules<'py>( view: &ArrayV, field: &Field, py: Python<'py>, -) -> PyResult<(PyObject, PyObject)> { +) -> PyResult<(Py, Py)> { let schema = Schema::from(vec![field.clone()]); let array = Arc::new(view.array.clone()); let (arr_ptr, sch_ptr) = @@ -882,7 +882,7 @@ pub fn array_view_to_capsules<'py>( pub fn table_view_to_stream_capsule<'py>( view: &TableV, py: Python<'py>, -) -> PyResult { +) -> PyResult> { let fields: Vec = view.fields.iter().map(|f| (**f).clone()).collect(); let metadata = if view.name.is_empty() { @@ -926,7 +926,7 @@ pub fn table_view_to_stream_capsule<'py>( pub fn super_table_view_to_stream_capsule<'py>( view: &SuperTableV, py: Python<'py>, -) -> PyResult { +) -> PyResult> { if view.slices.is_empty() { return Err(pyo3::exceptions::PyValueError::new_err( "Cannot export empty SuperTableV as stream capsule", @@ -975,7 +975,7 @@ pub fn super_table_view_to_stream_capsule<'py>( pub fn super_array_view_to_stream_capsule<'py>( view: &SuperArrayV, py: Python<'py>, -) -> PyResult { +) -> PyResult> { if view.slices.is_empty() { return Err(pyo3::exceptions::PyValueError::new_err( "Cannot export empty SuperArrayV as stream capsule", @@ -1014,7 +1014,7 @@ pub fn super_array_view_to_stream_capsule<'py>( pub fn super_array_to_stream_capsule<'py>( super_array: &SuperArray, py: Python<'py>, -) -> PyResult { +) -> PyResult> { let chunks = super_array.chunks(); if chunks.is_empty() { return Err(pyo3::exceptions::PyValueError::new_err( diff --git a/pyo3/src/ffi/to_rust.rs b/pyo3/src/ffi/to_rust.rs index 5350e54..2d33ebc 100644 --- a/pyo3/src/ffi/to_rust.rs +++ b/pyo3/src/ffi/to_rust.rs @@ -117,7 +117,7 @@ fn import_capsule_array(obj: &Bound) -> PyMinarrowResult { PyMinarrowError::PyArrow(format!("Failed to call __arrow_c_array__: {}", e)) })?; - let tuple: &Bound = result.downcast().map_err(|e| { + let tuple: &Bound = result.cast().map_err(|e| { PyMinarrowError::PyArrow(format!( "__arrow_c_array__ did not return a tuple: {}", e diff --git a/pyo3/src/lib.rs b/pyo3/src/lib.rs index 49b0d06..2180f21 100644 --- a/pyo3/src/lib.rs +++ b/pyo3/src/lib.rs @@ -113,6 +113,9 @@ pub mod error; pub mod ffi; pub mod types; +#[cfg(feature = "ndarray")] +use crate::ffi::dlpack::PyNdArrayInner; + // Re-export the main types for ease of use pub use error::{PyMinarrowError, PyMinarrowResult}; pub use types::{ @@ -259,7 +262,7 @@ fn export_chunked_stream_capsule(py: Python, arr: PyChunkedArray) -> PyResult, + capsule: Option>, } #[pymethods] @@ -269,7 +272,7 @@ impl ArrowStream { /// Returns the underlying ArrowArrayStream capsule. The capsule can /// only be consumed once - subsequent calls raise ValueError. #[pyo3(signature = (requested_schema=None))] - fn __arrow_c_stream__(&mut self, requested_schema: Option) -> PyResult { + fn __arrow_c_stream__(&mut self, requested_schema: Option>) -> PyResult> { let _ = requested_schema; self.capsule.take().ok_or_else(|| { pyo3::exceptions::PyValueError::new_err( @@ -285,8 +288,8 @@ impl ArrowStream { /// e.g. `pa.array(obj)`. #[pyclass(name = "ArrowArray")] struct ArrowArrayWrapper { - schema_capsule: Option, - array_capsule: Option, + schema_capsule: Option>, + array_capsule: Option>, } #[pymethods] @@ -298,8 +301,8 @@ impl ArrowArrayWrapper { #[pyo3(signature = (requested_schema=None))] fn __arrow_c_array__( &mut self, - requested_schema: Option, - ) -> PyResult<(PyObject, PyObject)> { + requested_schema: Option>, + ) -> PyResult<(Py, Py)> { let _ = requested_schema; let schema = self.schema_capsule.take(); let array = self.array_capsule.take(); @@ -312,6 +315,187 @@ impl ArrowArrayWrapper { } } +/// N-dimensional f32/f64 tensor with zero-copy DLPack interchange. +/// +/// The Rust-produces / Python-consumes counterpart of minarrow-py's +/// `NdArray` - construct it from a minarrow [`NdArray`](minarrow::NdArray) +/// via `From`, hand it to NumPy or PyTorch through `__dlpack__`, and take +/// tensors back with `from_dlpack`. Building tensors from Python +/// sequences is minarrow-py's job. +#[cfg(feature = "ndarray")] +#[pyclass(name = "NdArray")] +pub struct PyNdArray(pub PyNdArrayInner); + +#[cfg(feature = "ndarray")] +impl From> for PyNdArray { + fn from(ndarray: minarrow::NdArray) -> Self { + PyNdArray(PyNdArrayInner::from(ndarray)) + } +} + +#[cfg(feature = "ndarray")] +impl From> for PyNdArray { + fn from(ndarray: minarrow::NdArray) -> Self { + PyNdArray(PyNdArrayInner::from(ndarray)) + } +} + +#[cfg(feature = "ndarray")] +impl From> for PyNdArray { + fn from(view: minarrow::NdArrayV) -> Self { + PyNdArray(PyNdArrayInner::from(view)) + } +} + +#[cfg(feature = "ndarray")] +impl From> for PyNdArray { + fn from(view: minarrow::NdArrayV) -> Self { + PyNdArray(PyNdArrayInner::from(view)) + } +} + +#[cfg(feature = "ndarray")] +#[pymethods] +impl PyNdArray { + /// Dimension sizes. + #[getter] + fn shape<'py>(&self, py: Python<'py>) -> PyResult> { + match &self.0 { + PyNdArrayInner::F32(a) => pyo3::types::PyTuple::new(py, a.shape()), + PyNdArrayInner::F64(a) => pyo3::types::PyTuple::new(py, a.shape()), + PyNdArrayInner::F32View(v) => pyo3::types::PyTuple::new(py, v.shape()), + PyNdArrayInner::F64View(v) => pyo3::types::PyTuple::new(py, v.shape()), + } + } + + /// Element strides per dimension. + #[getter] + fn strides<'py>(&self, py: Python<'py>) -> PyResult> { + match &self.0 { + PyNdArrayInner::F32(a) => pyo3::types::PyTuple::new(py, a.strides()), + PyNdArrayInner::F64(a) => pyo3::types::PyTuple::new(py, a.strides()), + PyNdArrayInner::F32View(v) => pyo3::types::PyTuple::new(py, v.strides()), + PyNdArrayInner::F64View(v) => pyo3::types::PyTuple::new(py, v.strides()), + } + } + + /// Number of dimensions. + #[getter] + fn ndim(&self) -> usize { + match &self.0 { + PyNdArrayInner::F32(a) => a.ndim(), + PyNdArrayInner::F64(a) => a.ndim(), + PyNdArrayInner::F32View(v) => v.ndim(), + PyNdArrayInner::F64View(v) => v.ndim(), + } + } + + /// Element type name, `float32` or `float64`. + #[getter] + fn dtype(&self) -> &'static str { + match &self.0 { + PyNdArrayInner::F32(_) => "float32", + PyNdArrayInner::F64(_) => "float64", + PyNdArrayInner::F32View(_) => "float32", + PyNdArrayInner::F64View(_) => "float64", + } + } + + /// Total element count. + #[getter] + fn size(&self) -> usize { + match &self.0 { + PyNdArrayInner::F32(a) => a.len(), + PyNdArrayInner::F64(a) => a.len(), + PyNdArrayInner::F32View(v) => v.len(), + PyNdArrayInner::F64View(v) => v.len(), + } + } + + fn __len__(&self) -> PyResult { + let shape = match &self.0 { + PyNdArrayInner::F32(a) => a.shape(), + PyNdArrayInner::F64(a) => a.shape(), + PyNdArrayInner::F32View(v) => v.shape(), + PyNdArrayInner::F64View(v) => v.shape(), + }; + shape + .first() + .copied() + .ok_or_else(|| pyo3::exceptions::PyTypeError::new_err("len() of a rank-zero NdArray")) + } + + fn __repr__(&self) -> String { + let shape = match &self.0 { + PyNdArrayInner::F32(a) => a.shape().to_vec(), + PyNdArrayInner::F64(a) => a.shape().to_vec(), + PyNdArrayInner::F32View(v) => v.shape().to_vec(), + PyNdArrayInner::F64View(v) => v.shape().to_vec(), + }; + format!("NdArray(shape={:?}, dtype={})", shape, self.dtype()) + } + + /// DLPack producer entry point. Returns a capsule the consumer owns. + /// See [`ffi::dlpack::export_dlpack`] for the ABI and copy semantics. + #[pyo3(signature = (*, stream=None, max_version=None, dl_device=None, copy=None))] + fn __dlpack__( + &self, + py: Python<'_>, + stream: Option<&Bound<'_, PyAny>>, + max_version: Option<(u32, u32)>, + dl_device: Option<(i32, i32)>, + copy: Option, + ) -> PyResult> { + ffi::dlpack::export_dlpack(py, &self.0, stream, max_version, dl_device, copy) + } + + /// DLPack device query. Minarrow data lives on the CPU. + fn __dlpack_device__(&self) -> (i32, i32) { + (1, 0) + } + + /// Import from a compatible CPU f32/f64 DLPack producer, such as a NumPy + /// or PyTorch tensor, or from a raw DLPack capsule. A suitably aligned + /// buffer can be shared; otherwise the data copies into an aligned buffer. + #[staticmethod] + fn from_dlpack(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult { + ffi::dlpack::import_dlpack(py, obj).map(PyNdArray) + } + + /// Hand to NumPy as an `ndarray` via the capsule protocol. + fn to_numpy(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let numpy = py.import("numpy")?; + Ok(numpy.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to PyTorch as a `torch.Tensor` via the capsule protocol. + fn to_pytorch(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let torch = py.import("torch")?; + Ok(torch.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to JAX as a `jax.Array` via the capsule protocol. + fn to_jax(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let jax_numpy = py.import("jax.numpy")?; + Ok(jax_numpy.call_method1("from_dlpack", (slf,))?.unbind()) + } + + /// Hand to TensorFlow as a `tf.Tensor`. TensorFlow's DLPack entry + /// takes the capsule itself rather than the producer object. + fn to_tensorflow(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let capsule = slf.borrow().__dlpack__(py, None, None, None, None)?; + let dlpack = py.import("tensorflow.experimental.dlpack")?; + Ok(dlpack.call_method1("from_dlpack", (capsule,))?.unbind()) + } + + /// Hand to CuPy via the capsule protocol. CuPy holds device memory, + /// so this copies host data to the GPU on import. + fn to_cupy(slf: Bound<'_, Self>, py: Python<'_>) -> PyResult> { + let cupy = py.import("cupy")?; + Ok(cupy.call_method1("from_dlpack", (slf,))?.unbind()) + } +} + // Data generators - these produce protocol-conforming objects /// Generate sample data entirely in Rust and return as an ArrowStream. @@ -416,6 +600,8 @@ fn minarrow_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { // PyCapsule protocol wrapper types m.add_class::()?; m.add_class::()?; + #[cfg(feature = "ndarray")] + m.add_class::()?; // Data generators (return protocol-conforming objects) m.add_function(wrap_pyfunction!(generate_sample_batch, m)?)?; diff --git a/pyo3/src/types.rs b/pyo3/src/types.rs index dd00e80..5f67bb7 100644 --- a/pyo3/src/types.rs +++ b/pyo3/src/types.rs @@ -22,6 +22,7 @@ use minarrow::{ TableV, }; use pyo3::prelude::*; +use pyo3::Borrowed; use std::sync::Arc; use crate::ffi::{to_py, to_rust}; @@ -114,9 +115,11 @@ impl AsRef for PyArray { } } -impl<'py> FromPyObject<'py> for PyArray { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let field_array = to_rust::array_to_rust(ob)?; +impl<'py> FromPyObject<'_, 'py> for PyArray { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult { + let field_array = to_rust::array_to_rust(&ob)?; Ok(PyArray(field_array)) } } @@ -190,9 +193,11 @@ impl AsRef for PyRecordBatch { } } -impl<'py> FromPyObject<'py> for PyRecordBatch { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let table = to_rust::record_batch_to_rust(ob)?; +impl<'py> FromPyObject<'_, 'py> for PyRecordBatch { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult { + let table = to_rust::record_batch_to_rust(&ob)?; Ok(PyRecordBatch(table)) } } @@ -309,9 +314,11 @@ impl AsRef for PyTable { } } -impl<'py> FromPyObject<'py> for PyTable { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let table = to_rust::table_to_rust(ob)?; +impl<'py> FromPyObject<'_, 'py> for PyTable { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult { + let table = to_rust::table_to_rust(&ob)?; Ok(PyTable(table)) } } @@ -384,9 +391,11 @@ impl AsRef for PyChunkedArray { } } -impl<'py> FromPyObject<'py> for PyChunkedArray { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let array = to_rust::chunked_array_to_rust(ob)?; +impl<'py> FromPyObject<'_, 'py> for PyChunkedArray { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult { + let array = to_rust::chunked_array_to_rust(&ob)?; Ok(PyChunkedArray(array)) } } diff --git a/pyo3/tests/atomic_tests.rs b/pyo3/tests/atomic_tests.rs index 3071652..37350d8 100644 --- a/pyo3/tests/atomic_tests.rs +++ b/pyo3/tests/atomic_tests.rs @@ -65,9 +65,9 @@ use pyo3::types::IntoPyDict; use std::sync::Arc; fn main() -> PyResult<()> { - pyo3::prepare_freethreaded_python(); + Python::initialize(); - Python::with_gil(|py| { + Python::attach(|py| { println!("=== MinArrow <-> PyArrow Comprehensive Tests ===\n"); let mut passed = 0; diff --git a/pyo3/tests/python_roundtrip.rs b/pyo3/tests/python_roundtrip.rs index 5331d10..be86018 100644 --- a/pyo3/tests/python_roundtrip.rs +++ b/pyo3/tests/python_roundtrip.rs @@ -35,10 +35,9 @@ use minarrow_pyo3::{PyArray, PyRecordBatch}; use pyo3::prelude::*; fn main() -> PyResult<()> { - // Initialise Python - pyo3::prepare_freethreaded_python(); + Python::initialize(); - Python::with_gil(|py| { + Python::attach(|py| { println!("=== MinArrow <-> PyArrow Roundtrip Example ===\n"); // Test 1: Integer Array roundtrip diff --git a/src/aliases.rs b/src/aliases.rs index d349de8..33a54be 100644 --- a/src/aliases.rs +++ b/src/aliases.rs @@ -56,6 +56,13 @@ use crate::{ #[cfg(feature = "chunked")] use crate::SuperTable; +#[cfg(feature = "ndarray")] +use crate::NdArray; +#[cfg(all(feature = "ndarray", feature = "chunked"))] +use crate::SuperNdArray; +#[cfg(feature = "xarray")] +use crate::XArray; + /// # RecordBatch /// /// Standard Arrow `Record Batch`. Alias of *Minarrow* `Table`. @@ -150,6 +157,22 @@ pub type ArrayVT<'a> = (&'a Array, Offset, Length); /// Windowed ***V**iew **T**uple* for Bitmask pub type BitmaskVT<'a> = (&'a Bitmask, Offset, Length); +/// Windowed ***V**iew **T**uple* for NdArray. Offset and length are +/// axis-0 observation counts. +#[cfg(feature = "ndarray")] +pub type NdArrayVT<'a, T> = (&'a NdArray, Offset, Length); + +/// Windowed ***V**iew **T**uple* for SuperNdArray. Offset and length are +/// axis-0 observation counts across batches. For chunk-spanning access +/// with methods, use [`SuperNdArrayV`](crate::SuperNdArrayV). +#[cfg(all(feature = "ndarray", feature = "chunked"))] +pub type SuperNdArrayVT<'a, T> = (&'a SuperNdArray, Offset, Length); + +/// Windowed ***V**iew **T**uple* for XArray. Offset and length are +/// axis-0 observation counts. +#[cfg(feature = "xarray")] +pub type XArrayVT<'a, T> = (&'a XArray, Offset, Length); + /// Subset per respective table within the cube /// /// Respects the means in which each table is windowed, diff --git a/src/enums/scalar.rs b/src/enums/scalar.rs index b18cca0..dfac718 100644 --- a/src/enums/scalar.rs +++ b/src/enums/scalar.rs @@ -22,6 +22,8 @@ //! to match to one of a range of possible values. use std::convert::From; +#[cfg(feature = "scalar_type")] +use std::fmt::{self, Display, Formatter}; #[cfg(feature = "scalar_type")] use num_traits::Pow; @@ -43,7 +45,7 @@ use crate::{Array, Bitmask, BooleanArray, FloatArray, IntegerArray, MaskedArray, /// - There are also `try_` methods that can be used to attempt it gracefully /// without the risk of panicking. #[cfg(feature = "scalar_type")] -#[derive(Debug, Clone, Display, Default)] +#[derive(Debug, Clone, Default)] pub enum Scalar { #[default] Null, @@ -77,6 +79,39 @@ pub enum Scalar { Interval, } +#[cfg(feature = "scalar_type")] +impl Display for Scalar { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Scalar::Null => f.write_str("Null"), + Scalar::Boolean(v) => Display::fmt(v, f), + #[cfg(feature = "extended_numeric_types")] + Scalar::Int8(v) => Display::fmt(v, f), + #[cfg(feature = "extended_numeric_types")] + Scalar::Int16(v) => Display::fmt(v, f), + Scalar::Int32(v) => Display::fmt(v, f), + Scalar::Int64(v) => Display::fmt(v, f), + #[cfg(feature = "extended_numeric_types")] + Scalar::UInt8(v) => Display::fmt(v, f), + #[cfg(feature = "extended_numeric_types")] + Scalar::UInt16(v) => Display::fmt(v, f), + Scalar::UInt32(v) => Display::fmt(v, f), + Scalar::UInt64(v) => Display::fmt(v, f), + Scalar::Float32(v) => Display::fmt(v, f), + Scalar::Float64(v) => Display::fmt(v, f), + Scalar::String32(v) => f.write_str(v), + #[cfg(feature = "large_string")] + Scalar::String64(v) => f.write_str(v), + #[cfg(feature = "datetime")] + Scalar::Datetime32(v) => Display::fmt(v, f), + #[cfg(feature = "datetime")] + Scalar::Datetime64(v) => Display::fmt(v, f), + #[cfg(feature = "datetime")] + Scalar::Interval => f.write_str("Interval"), + } + } +} + #[cfg(feature = "scalar_type")] impl Scalar { /// Casts the value to a bool @@ -2101,6 +2136,25 @@ mod tests { use super::*; use crate::Scalar::{Float32, Float64, Int32, String32, String64}; + #[test] + fn display() { + assert_eq!(Scalar::Null.to_string(), "Null"); + assert_eq!(Scalar::Boolean(true).to_string(), "true"); + assert_eq!(Scalar::Int32(-42).to_string(), "-42"); + assert_eq!(Scalar::Float64(1.5).to_string(), "1.5"); + assert_eq!(Scalar::String32("value".into()).to_string(), "value"); + + #[cfg(feature = "large_string")] + assert_eq!(Scalar::String64("large".into()).to_string(), "large"); + + #[cfg(feature = "datetime")] + { + assert_eq!(Scalar::Datetime32(123).to_string(), "123"); + assert_eq!(Scalar::Datetime64(456).to_string(), "456"); + assert_eq!(Scalar::Interval.to_string(), "Interval"); + } + } + #[test] fn test_bool() { assert_eq!(Scalar::Boolean(true).bool(), true); diff --git a/src/enums/value/conversions.rs b/src/enums/value/conversions.rs index afa0fae..566be54 100644 --- a/src/enums/value/conversions.rs +++ b/src/enums/value/conversions.rs @@ -22,8 +22,18 @@ use super::impls::value_variant_name; use crate::Cube; #[cfg(feature = "matrix")] use crate::Matrix; +#[cfg(feature = "ndarray")] +use crate::NdArray; +#[cfg(all(feature = "ndarray", feature = "views"))] +use crate::NdArrayV; #[cfg(feature = "scalar_type")] use crate::Scalar; +#[cfg(all(feature = "ndarray", feature = "chunked"))] +use crate::SuperNdArray; +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +use crate::SuperNdArrayV; +#[cfg(feature = "xarray")] +use crate::XArray; use crate::traits::concatenate::Concatenate; use crate::traits::masked_array::MaskedArray; use crate::{Array, FieldArray, Table, enums::error::MinarrowError}; @@ -895,6 +905,46 @@ impl From for Value { } } +#[cfg(feature = "ndarray")] +impl From> for Value { + #[inline] + fn from(v: NdArray) -> Self { + Value::NdArray(Arc::new(v)) + } +} + +#[cfg(all(feature = "ndarray", feature = "views"))] +impl From> for Value { + #[inline] + fn from(v: NdArrayV) -> Self { + Value::NdArrayView(Arc::new(v)) + } +} + +#[cfg(all(feature = "ndarray", feature = "chunked"))] +impl From> for Value { + #[inline] + fn from(v: SuperNdArray) -> Self { + Value::SuperNdArray(Arc::new(v)) + } +} + +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +impl From> for Value { + #[inline] + fn from(v: SuperNdArrayV) -> Self { + Value::SuperNdArrayView(Arc::new(v)) + } +} + +#[cfg(feature = "xarray")] +impl From> for Value { + #[inline] + fn from(v: XArray) -> Self { + Value::XArray(Arc::new(v)) + } +} + /// Wrap a `NumericArrayV` as `Value::ArrayView`. The view's offset and length /// are preserved; the inner `NumericArray` flows through `Array::NumericArray`. #[cfg(feature = "views")] @@ -996,6 +1046,19 @@ impl TryFrom for Array { Value::SuperTableView(_) => Err(err()), #[cfg(feature = "matrix")] Value::Matrix(_) => Err(err()), + #[cfg(feature = "ndarray")] + Value::NdArray(inner) => { + let nd = Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone()); + nd.to_array() + } + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(_) => Err(err()), + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(_) => Err(err()), + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(_) => Err(err()), + #[cfg(feature = "xarray")] + Value::XArray(_) => Err(err()), #[cfg(feature = "cube")] Value::Cube(_) => Err(err()), Value::VecValue(inner) => { @@ -1302,6 +1365,22 @@ impl TryFrom for Table { Value::SuperTableView(_) => Err(err()), #[cfg(feature = "matrix")] Value::Matrix(_) => Err(err()), + #[cfg(feature = "ndarray")] + Value::NdArray(inner) => { + let nd = Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone()); + nd.to_table(None) + } + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(_) => Err(err()), + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(_) => Err(err()), + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(_) => Err(err()), + #[cfg(feature = "xarray")] + Value::XArray(inner) => { + let xa = Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone()); + xa.to_table() + } #[cfg(feature = "cube")] Value::Cube(_) => Err(err()), // VecValue terminal coercion: engine wire produced by `Long` fan-out gather. @@ -1520,6 +1599,91 @@ impl TryFrom for SuperTableV { } } +#[cfg(feature = "ndarray")] +impl TryFrom for NdArray { + type Error = MinarrowError; + fn try_from(v: Value) -> Result { + match v { + Value::NdArray(inner) => { + Ok(Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone())) + } + _ => Err(MinarrowError::TypeError { + from: "Value", + to: "NdArray", + message: Some("Value type mismatch".to_owned()), + }), + } + } +} + +#[cfg(all(feature = "ndarray", feature = "views"))] +impl TryFrom for NdArrayV { + type Error = MinarrowError; + fn try_from(v: Value) -> Result { + match v { + Value::NdArrayView(inner) => { + Ok(Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone())) + } + _ => Err(MinarrowError::TypeError { + from: "Value", + to: "NdArrayV", + message: Some("Value type mismatch".to_owned()), + }), + } + } +} + +#[cfg(all(feature = "ndarray", feature = "chunked"))] +impl TryFrom for SuperNdArray { + type Error = MinarrowError; + fn try_from(v: Value) -> Result { + match v { + Value::SuperNdArray(inner) => { + Ok(Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone())) + } + _ => Err(MinarrowError::TypeError { + from: "Value", + to: "SuperNdArray", + message: Some("Value type mismatch".to_owned()), + }), + } + } +} + +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +impl TryFrom for SuperNdArrayV { + type Error = MinarrowError; + fn try_from(v: Value) -> Result { + match v { + Value::SuperNdArrayView(inner) => { + Ok(Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone())) + } + _ => Err(MinarrowError::TypeError { + from: "Value", + to: "SuperNdArrayV", + message: Some("Value type mismatch".to_owned()), + }), + } + } +} + +#[cfg(feature = "xarray")] +impl TryFrom for XArray { + type Error = MinarrowError; + fn try_from(v: Value) -> Result { + match v { + Value::XArray(inner) => { + Ok(Arc::try_unwrap(inner).unwrap_or_else(|arc| (*arc).clone())) + } + _ => Err(MinarrowError::TypeError { + from: "Value", + to: "XArray", + message: Some("Value type mismatch".to_owned()), + }), + } + } +} + // Recursive Container Conversions impl From> for Value { @@ -2100,6 +2264,48 @@ mod accessor_tests { let result = val.try_st(); assert!(result.is_ok()); } + + #[cfg(feature = "ndarray")] + #[test] + fn test_ndarray_value_roundtrips() { + let nd = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + + let back = NdArray::::try_from(Value::from(nd.clone())).unwrap(); + assert_eq!(back, nd); + assert!(NdArray::::try_from(Value::Table(Arc::new(Table::new_empty()))).is_err()); + + #[cfg(feature = "views")] + { + let view = nd.as_view(); + let back = NdArrayV::::try_from(Value::from(view)).unwrap(); + assert_eq!(back.to_ndarray(), nd); + assert!(NdArrayV::::try_from(Value::from(nd.clone())).is_err()); + } + + #[cfg(feature = "chunked")] + { + let snd = SuperNdArray::from_batches(vec![nd.clone(), nd.clone()], "s"); + let back = SuperNdArray::::try_from(Value::from(snd.clone())).unwrap(); + assert_eq!(back, snd); + assert!(SuperNdArray::::try_from(Value::from(nd.clone())).is_err()); + + #[cfg(feature = "views")] + { + let sv = snd.slice(0, 4); + let back = SuperNdArrayV::::try_from(Value::from(sv.clone())).unwrap(); + assert_eq!(back, sv); + assert!(SuperNdArrayV::::try_from(Value::from(snd.clone())).is_err()); + } + } + + #[cfg(feature = "xarray")] + { + let xa = XArray::new(nd.clone(), &["obs", "feat"]); + let back = XArray::::try_from(Value::from(xa.clone())).unwrap(); + assert_eq!(back, xa); + assert!(XArray::::try_from(Value::from(nd)).is_err()); + } + } } // VecValue -> typed-T terminal coercion tests. Exercises the engine wire format diff --git a/src/enums/value/impls.rs b/src/enums/value/impls.rs index 823639d..68b9a62 100644 --- a/src/enums/value/impls.rs +++ b/src/enums/value/impls.rs @@ -47,6 +47,16 @@ impl PartialEq for Value { (FieldArray(a), FieldArray(b)) => **a == **b, #[cfg(feature = "matrix")] (Matrix(a), Matrix(b)) => a == b, + #[cfg(feature = "ndarray")] + (NdArray(a), NdArray(b)) => **a == **b, + #[cfg(all(feature = "ndarray", feature = "views"))] + (NdArrayView(a), NdArrayView(b)) => **a == **b, + #[cfg(all(feature = "ndarray", feature = "chunked"))] + (SuperNdArray(a), SuperNdArray(b)) => **a == **b, + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + (SuperNdArrayView(a), SuperNdArrayView(b)) => **a == **b, + #[cfg(feature = "xarray")] + (XArray(a), XArray(b)) => **a == **b, #[cfg(feature = "cube")] (Cube(a), Cube(b)) => **a == **b, (Custom(a), Custom(b)) => a.eq_box(&**b), @@ -99,6 +109,16 @@ impl Shape for Value { Value::FieldArray(field_array) => field_array.shape(), #[cfg(feature = "matrix")] Value::Matrix(matrix) => matrix.shape(), + #[cfg(feature = "ndarray")] + Value::NdArray(nd) => Shape::shape(nd.as_ref()), + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(v) => Shape::shape(v.as_ref()), + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(snd) => Shape::shape(snd.as_ref()), + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(v) => Shape::shape(v.as_ref()), + #[cfg(feature = "xarray")] + Value::XArray(xa) => Shape::shape(xa.as_ref()), #[cfg(feature = "cube")] Value::Cube(cube) => cube.shape(), Value::VecValue(vec_value) => { @@ -265,6 +285,47 @@ impl Concatenate for Value { Ok(Value::Cube(Arc::new(a.concat(b)?))) } + // NdArray + NdArray -> NdArray + #[cfg(feature = "ndarray")] + (NdArray(a), NdArray(b)) => { + let a = Arc::try_unwrap(a).unwrap_or_else(|arc| (*arc).clone()); + let b = Arc::try_unwrap(b).unwrap_or_else(|arc| (*arc).clone()); + Ok(Value::NdArray(Arc::new(a.concat(b)?))) + } + + // NdArrayView + NdArrayView -> NdArray, materialising both + // windows, mirroring the broadcast semantics for view pairs. + #[cfg(all(feature = "ndarray", feature = "views"))] + (NdArrayView(a), NdArrayView(b)) => { + let a = a.to_ndarray(); + let b = b.to_ndarray(); + Ok(Value::NdArray(Arc::new(a.concat(b)?))) + } + + // SuperNdArray + SuperNdArray -> SuperNdArray + #[cfg(all(feature = "ndarray", feature = "chunked"))] + (SuperNdArray(a), SuperNdArray(b)) => { + let a = Arc::try_unwrap(a).unwrap_or_else(|arc| (*arc).clone()); + let b = Arc::try_unwrap(b).unwrap_or_else(|arc| (*arc).clone()); + Ok(Value::SuperNdArray(Arc::new(a.concat(b)?))) + } + + // SuperNdArrayView + SuperNdArrayView -> SuperNdArrayView, zero-copy + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + (SuperNdArrayView(a), SuperNdArrayView(b)) => { + let a = Arc::try_unwrap(a).unwrap_or_else(|arc| (*arc).clone()); + let b = Arc::try_unwrap(b).unwrap_or_else(|arc| (*arc).clone()); + Ok(Value::SuperNdArrayView(Arc::new(a.concat(b)?))) + } + + // XArray + XArray -> XArray + #[cfg(feature = "xarray")] + (XArray(a), XArray(b)) => { + let a = Arc::try_unwrap(a).unwrap_or_else(|arc| (*arc).clone()); + let b = Arc::try_unwrap(b).unwrap_or_else(|arc| (*arc).clone()); + Ok(Value::XArray(Arc::new(a.concat(b)?))) + } + // Chunked types #[cfg(feature = "chunked")] (SuperArray(a), SuperArray(b)) => { @@ -490,6 +551,16 @@ pub(crate) fn value_variant_name(value: &Value) -> &'static str { Value::FieldArray(_) => "FieldArray", #[cfg(feature = "matrix")] Value::Matrix(_) => "Matrix", + #[cfg(feature = "ndarray")] + Value::NdArray(_) => "NdArray", + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(_) => "NdArrayView", + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(_) => "SuperNdArray", + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(_) => "SuperNdArrayView", + #[cfg(feature = "xarray")] + Value::XArray(_) => "XArray", #[cfg(feature = "cube")] Value::Cube(_) => "Cube", Value::VecValue(_) => "VecValue", @@ -506,8 +577,10 @@ pub(crate) fn value_variant_name(value: &Value) -> &'static str { // Consolidate for Vec +#[cfg(feature = "chunked")] use crate::traits::consolidate::Consolidate; +#[cfg(feature = "chunked")] impl Consolidate for Vec { type Output = Value; @@ -756,4 +829,129 @@ mod concat_tests { panic!("Expected IncompatibleTypeError"); } } + + #[cfg(feature = "ndarray")] + #[test] + fn test_value_len_ndarray_types() { + use crate::NdArray; + + // Column-major [3, 2] with 3 axis-0 observations + let nd = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(Value::NdArray(Arc::new(nd.clone())).len(), 3); + + #[cfg(feature = "views")] + assert_eq!(Value::NdArrayView(Arc::new(nd.as_view())).len(), 3); + + #[cfg(feature = "chunked")] + { + use crate::SuperNdArray; + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0], &[1, 2]), + ], + "s", + ); + assert_eq!(Value::SuperNdArray(Arc::new(snd.clone())).len(), 3); + #[cfg(feature = "views")] + assert_eq!(Value::SuperNdArrayView(Arc::new(snd.slice(0, 2))).len(), 2); + } + + #[cfg(feature = "xarray")] + { + use crate::XArray; + let xa = XArray::new(nd, &["obs", "feat"]); + assert_eq!(Value::XArray(Arc::new(xa)).len(), 3); + } + } + + #[cfg(all(feature = "ndarray", feature = "views"))] + #[test] + fn test_value_slice_ndarray() { + use crate::NdArray; + + let nd = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = Value::NdArray(Arc::new(nd)); + + // The full window keeps the shape. + let full = v.slice(0, v.len()); + let Value::NdArrayView(view) = full else { + panic!("Expected Value::NdArrayView"); + }; + assert_eq!(view.shape(), &[3, 2]); + assert_eq!(view.get(&[0, 0]), 1.0); + assert_eq!(view.get(&[2, 1]), 6.0); + + // A partial window covers rows 1..3. + let window = v.slice(1, 2); + let Value::NdArrayView(view) = window else { + panic!("Expected Value::NdArrayView"); + }; + assert_eq!(view.shape(), &[2, 2]); + assert_eq!(view.get(&[0, 0]), 2.0); + assert_eq!(view.get(&[1, 0]), 3.0); + assert_eq!(view.get(&[0, 1]), 5.0); + assert_eq!(view.get(&[1, 1]), 6.0); + } + + #[cfg(all(feature = "ndarray", feature = "views"))] + #[test] + #[should_panic(expected = "out of bounds")] + fn test_value_slice_ndarray_out_of_bounds_panics() { + use crate::NdArray; + + let nd = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = Value::NdArray(Arc::new(nd)); + let _ = v.slice(2, 2); + } + + #[cfg(all(feature = "xarray", feature = "views", feature = "select"))] + #[test] + fn test_value_slice_xarray_narrows_coords() { + use crate::{FloatArray, NdArray, NumericArray, XArray}; + + let mut xa = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]), + &["obs", "feat"], + ); + xa.assign_coords( + "obs", + Array::NumericArray(NumericArray::Float64(Arc::new( + FloatArray::from_slice(&[10.0, 20.0, 30.0]), + ))), + ); + let v = Value::XArray(Arc::new(xa)); + + let sliced = v.slice(1, 2); + let Value::XArray(out) = sliced else { + panic!("Expected Value::XArray"); + }; + assert_eq!(out.shape(), vec![2, 2]); + // The axis-0 coords narrow alongside the data window. + let coords = out.ax("obs").coords.as_ref().unwrap(); + let expected = Array::NumericArray(NumericArray::Float64(Arc::new( + FloatArray::from_slice(&[20.0, 30.0]), + ))); + assert_eq!(*coords, expected); + } + + #[cfg(all(feature = "ndarray", feature = "views"))] + #[test] + fn test_value_concat_ndarray_views() { + use crate::NdArray; + + let a = NdArray::from_slice(&[1.0, 2.0], &[2]); + let b = NdArray::from_slice(&[3.0], &[1]); + let va = Value::NdArrayView(Arc::new(a.as_view())); + let vb = Value::NdArrayView(Arc::new(b.as_view())); + + let out = va.concat(vb).unwrap(); + let Value::NdArray(nd) = out else { + panic!("Expected Value::NdArray"); + }; + assert_eq!(nd.shape(), &[3]); + assert_eq!(nd.get(&[0]), 1.0); + assert_eq!(nd.get(&[1]), 2.0); + assert_eq!(nd.get(&[2]), 3.0); + } } diff --git a/src/enums/value/mod.rs b/src/enums/value/mod.rs index cb0c6db..5f42024 100644 --- a/src/enums/value/mod.rs +++ b/src/enums/value/mod.rs @@ -36,8 +36,18 @@ mod conversions; use crate::Cube; #[cfg(feature = "matrix")] use crate::Matrix; +#[cfg(feature = "ndarray")] +use crate::NdArray; +#[cfg(all(feature = "ndarray", feature = "views"))] +use crate::NdArrayV; #[cfg(feature = "scalar_type")] use crate::Scalar; +#[cfg(all(feature = "ndarray", feature = "chunked"))] +use crate::SuperNdArray; +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +use crate::SuperNdArrayV; +#[cfg(feature = "xarray")] +use crate::XArray; use crate::{Array, FieldArray, Table, traits::custom_value::CustomValue}; use std::sync::Arc; @@ -91,6 +101,16 @@ pub enum Value { SuperTableView(Arc), #[cfg(feature = "matrix")] Matrix(Arc), + #[cfg(feature = "ndarray")] + NdArray(Arc>), + #[cfg(all(feature = "ndarray", feature = "views"))] + NdArrayView(Arc>), + #[cfg(all(feature = "ndarray", feature = "chunked"))] + SuperNdArray(Arc>), + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + SuperNdArrayView(Arc>), + #[cfg(feature = "xarray")] + XArray(Arc>), #[cfg(feature = "cube")] Cube(Arc), VecValue(Arc>), @@ -123,7 +143,9 @@ impl Value { /// Computes the logical row/element count for the batch's input `Value`. /// /// This normalises the various `Value` representations so callers can consistently pass a - /// `[start, len)` range to `execute_fn`. + /// `[start, len)` range to `execute_fn`. For the n-dimensional types + /// this is the leading-axis observation count, matching the units + /// `slice` windows over. #[inline] pub fn len(&self) -> usize { match self { @@ -157,6 +179,21 @@ impl Value { #[cfg(feature = "matrix")] Value::Matrix(m) => m.len(), + #[cfg(feature = "ndarray")] + Value::NdArray(nd) => nd.shape()[0], + + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(v) => v.shape()[0], + + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(snd) => snd.n_obs(), + + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(v) => v.n_obs(), + + #[cfg(feature = "xarray")] + Value::XArray(xa) => xa.shape()[0], + #[cfg(feature = "cube")] Value::Cube(c) => c.len(), @@ -226,6 +263,56 @@ impl Value { #[cfg(feature = "matrix")] Value::Matrix(_) => unimplemented!("Matrix slicing"), + #[cfg(feature = "ndarray")] + Value::NdArray(nd) => { + assert!( + offset + length <= nd.shape()[0], + "Value::slice: window {}..{} out of bounds for axis 0 (size {})", + offset, offset + length, nd.shape()[0] + ); + let mut window_shape = vec![length]; + window_shape.extend_from_slice(&nd.shape()[1..]); + Value::NdArrayView(Arc::new(NdArrayV::new( + nd.as_ref().clone(), + offset * nd.strides()[0], + &window_shape, + nd.strides(), + ))) + } + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(v) => { + assert!( + offset + length <= v.shape()[0], + "Value::slice: window {}..{} out of bounds for axis 0 (size {})", + offset, offset + length, v.shape()[0] + ); + let mut window_shape = vec![length]; + window_shape.extend_from_slice(&v.shape()[1..]); + Value::NdArrayView(Arc::new(NdArrayV::new( + v.source.clone(), + v.offset + offset * v.strides()[0], + &window_shape, + v.strides(), + ))) + } + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(snd) => { + Value::SuperNdArrayView(Arc::new(snd.slice(offset, length))) + } + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(v) => { + Value::SuperNdArrayView(Arc::new(v.slice(offset, length))) + } + #[cfg(all(feature = "xarray", feature = "select"))] + Value::XArray(xa) => { + // An axis-0 window through select, which narrows the + // leading axis coords alongside the data. + let range = offset..offset + length; + let dim0 = xa.dim_names()[0].to_string(); + Value::XArray(Arc::new(xa.select(&[(dim0.as_str(), &range)]))) + } + #[cfg(all(feature = "xarray", not(feature = "select")))] + Value::XArray(_) => unimplemented!("XArray slicing requires the select feature"), #[cfg(feature = "cube")] Value::Cube(_) => unimplemented!("Cube slicing"), diff --git a/src/ffi/dlpack.rs b/src/ffi/dlpack.rs new file mode 100644 index 0000000..1573520 --- /dev/null +++ b/src/ffi/dlpack.rs @@ -0,0 +1,1330 @@ +//! # **DLPack FFI** - *Tensor interchange with Python and other runtimes* +//! +//! Implements the [DLPack](https://github.com/dmlc/dlpack) tensor interchange +//! standard for NdArray. Compatible ownership, alignment, and protocol choices +//! allow runtimes to share the allocation directly; the documented fallback +//! cases copy so that Minarrow's ownership and alignment guarantees still hold. +//! +//! ## Supported protocols +//! - **DLManagedTensor** (legacy, pre-1.0) - the capsule payload older +//! consumers expect +//! - **DLManagedTensorVersioned** (DLPack 1.x) - carries the spec version +//! and the flags word, including the read-only bit +//! - Export: [`export_to_dlpack`] / [`export_to_dlpack_versioned`] wrap an +//! NdArray, and [`export_view_to_dlpack`] / [`export_view_to_dlpack_versioned`] +//! wrap an NdArrayV window without copying +//! - Import: [`import_from_dlpack`] / [`import_from_dlpack_versioned`] wrap a +//! foreign managed tensor as an NdArray +//! +//! ## Version negotiation +//! A Python `__dlpack__(max_version=...)` implementation calls +//! [`export_to_dlpack_versioned`] when the consumer requests major version 1 +//! or above, and [`export_to_dlpack`] otherwise. The versioned capsule is +//! named `dltensor_versioned` and the legacy capsule `dltensor`. A consumed +//! capsule renames with a `used_` prefix per the protocol. +//! +//! ## Read-only flag +//! Versioned exports set `DLPACK_FLAG_BITMASK_READ_ONLY` when the backing +//! buffer is shared, since a consumer writing through the pointer would be +//! visible to every other reference. Imported read-only tensors land as +//! shared buffers, so Minarrow's copy-on-write semantics honour the flag - +//! the first mutation copies. +//! +//! ## Layout compatibility +//! NdArray's compact column-major layout is expressed through DLPack's +//! strides field, and the buffer is fully contiguous, so consumers receive +//! a dense tensor with no dead bytes. Consumers such as PyTorch, NumPy, JAX, +//! and TensorFlow decide whether to operate on those strides directly or +//! re-lay the tensor for a particular operation. +//! +//! ## Notes +//! - DLPack uses element strides, matching NdArray's convention +//! - DLPack has no null mask concept - NaN-based missing values align with this +//! - The allocation start is 64-byte aligned, matching the DLPack +//! recommendation for aligned data pointers +//! - Imported foreign buffers that are not 64-byte aligned copy into an +//! aligned `Vec64` so the crate's SIMD-alignment invariant holds + +use std::ffi::c_void; +use std::sync::Arc; + +use crate::enums::error::MinarrowError; +use crate::structs::buffer::Buffer; +use crate::structs::ndarray::NdArray; +use crate::structs::shared_buffer::SharedBuffer; +#[cfg(feature = "views")] +use crate::structs::views::ndarray_view::NdArrayV; +use crate::traits::type_unions::Float; + +// **************************************************************** +// DLPack C structs (matching the DLPack header) +// **************************************************************** + +/// The DLPack major version this crate produces and understands. +pub const DLPACK_MAJOR_VERSION: u32 = 1; + +/// The DLPack minor version this crate produces. +pub const DLPACK_MINOR_VERSION: u32 = 1; + +/// Flag bit marking the tensor data as read-only. +pub const DLPACK_FLAG_BITMASK_READ_ONLY: u64 = 1; + +/// Flag bit marking the tensor as a copy of the producer's data. +pub const DLPACK_FLAG_BITMASK_IS_COPIED: u64 = 2; + +/// Flag bit marking sub-byte element types as byte-padded. +pub const DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED: u64 = 4; + +/// Device type code from the DLPack spec. +/// +/// Held as a plain integer rather than a Rust enum, since a foreign +/// producer may send any code the spec defines, including ones added +/// after this crate was compiled. Reading an unknown discriminant into +/// a Rust enum would be undefined behaviour. +#[repr(transparent)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DLDeviceType(pub i32); + +impl DLDeviceType { + pub const CPU: Self = DLDeviceType(1); + pub const CUDA: Self = DLDeviceType(2); + pub const CUDA_HOST: Self = DLDeviceType(3); + pub const OPENCL: Self = DLDeviceType(4); + pub const VULKAN: Self = DLDeviceType(7); + pub const METAL: Self = DLDeviceType(8); + pub const VPI: Self = DLDeviceType(9); + pub const ROCM: Self = DLDeviceType(10); + pub const ROCM_HOST: Self = DLDeviceType(11); + pub const EXT_DEV: Self = DLDeviceType(12); + pub const CUDA_MANAGED: Self = DLDeviceType(13); + pub const ONE_API: Self = DLDeviceType(14); + pub const WEBGPU: Self = DLDeviceType(15); + pub const HEXAGON: Self = DLDeviceType(16); + pub const MAIA: Self = DLDeviceType(17); +} + +/// Device descriptor. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DLDevice { + pub device_type: DLDeviceType, + pub device_id: i32, +} + +impl DLDevice { + /// Host CPU device, id 0. The device every Minarrow buffer lives on, + /// and the value a Python `__dlpack_device__` implementation returns. + pub const fn cpu() -> Self { + DLDevice { device_type: DLDeviceType::CPU, device_id: 0 } + } +} + +/// Data type descriptor. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DLDataType { + /// Type code: 0=int, 1=uint, 2=float, 4=bfloat, 5=complex, 6=bool + pub code: u8, + /// Number of bits per element + pub bits: u8, + /// Number of SIMD lanes, typically 1 + pub lanes: u16, +} + +impl DLDataType { + /// f32 type descriptor. + pub const FLOAT32: Self = DLDataType { code: 2, bits: 32, lanes: 1 }; + + /// f64 type descriptor. + pub const FLOAT64: Self = DLDataType { code: 2, bits: 64, lanes: 1 }; + + /// Float descriptor for a Minarrow element type. + pub fn float() -> Self { + DLDataType { code: 2, bits: (std::mem::size_of::() * 8) as u8, lanes: 1 } + } +} + +/// DLPack spec version carried by the versioned managed tensor. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DLPackVersion { + pub major: u32, + pub minor: u32, +} + +/// The core DLPack tensor descriptor. Does not own the data. +#[repr(C)] +pub struct DLTensor { + /// Pointer to the start of the data buffer. + pub data: *mut c_void, + /// Device where the data resides. + pub device: DLDevice, + /// Number of dimensions. + pub ndim: i32, + /// Data type. + pub dtype: DLDataType, + /// Shape array with ndim elements. + pub shape: *mut i64, + /// Strides array with ndim elements, or null for C-contiguous. + pub strides: *mut i64, + /// Byte offset from data pointer to the start of actual data. + pub byte_offset: u64, +} + +/// DLPack legacy managed tensor with ownership semantics. +/// The deleter is called when the consumer is done with the tensor. +#[repr(C)] +pub struct DLManagedTensor { + pub dl_tensor: DLTensor, + pub manager_ctx: *mut c_void, + pub deleter: Option, +} + +unsafe impl Send for DLManagedTensor {} +unsafe impl Sync for DLManagedTensor {} + +/// DLPack 1.x managed tensor. Extends the legacy layout with the spec +/// version and a flags word. Field order is ABI - version leads so a +/// consumer can check compatibility before reading the rest. +#[repr(C)] +pub struct DLManagedTensorVersioned { + pub version: DLPackVersion, + pub manager_ctx: *mut c_void, + pub deleter: Option, + pub flags: u64, + pub dl_tensor: DLTensor, +} + +unsafe impl Send for DLManagedTensorVersioned {} +unsafe impl Sync for DLManagedTensorVersioned {} + +// **************************************************************** +// Export: NdArray / NdArrayV -> managed tensor +// **************************************************************** + +/// Keeps the exported source and the i64 shape/strides arrays alive for +/// the lifetime of the managed tensor. +struct DLPackHolder { + _source: A, + shape_i64: Vec, + strides_i64: Vec, +} + +/// Export an NdArray as a legacy DLPack managed tensor. +/// +/// Legacy `DLManagedTensor` carries no read-only flag, so consumers +/// treat the pointer as writable. A uniquely owned Minarrow allocation +/// transfers zero-copy. Shared or aliased storage copies first so a +/// consumer cannot mutate another reference through the legacy protocol. +/// Consumers on the 1.x protocol should prefer `to_dlpack_versioned`, +/// which can share the buffer and carry its read-only state. A buffer +/// without a stable owned allocation also copies first, since the consumer +/// holds the base pointer for the tensor's whole lifetime. +/// +/// Returns a [`DLPackTensor`] that manages the lifecycle. The NdArray's +/// buffer stays alive through its internal shared reference count. When +/// the owned tensor is dropped, the backing data is released. +/// +/// For FFI handoff (e.g. creating a PyCapsule), call `.into_raw()` to +/// transfer ownership to the consumer. +pub fn export_to_dlpack(ndarray: NdArray) -> DLPackTensor { + let ndarray = if ndarray.data.is_owned() && Arc::strong_count(&ndarray.data) == 1 { + ndarray + } else { + NdArray::from_buffer(ndarray.data.to_owned_copy(), ndarray.shape(), ndarray.strides()) + }; + let shape = ndarray.shape().to_vec(); + let strides = ndarray.strides().to_vec(); + let data_ptr = ndarray.as_slice().as_ptr() as *mut c_void; + DLPackTensor::from_parts(ndarray, data_ptr, 0, DLDataType::float::(), &shape, &strides) +} + +/// Export an NdArray as a DLPack 1.x versioned managed tensor. +/// +/// Sets `DLPACK_FLAG_BITMASK_READ_ONLY` when the backing buffer is +/// shared, since a consumer writing through the pointer would be visible +/// to every other reference. A uniquely owned buffer exports writable. +/// A buffer without a stable owned or shared allocation copies into an +/// owned one first, since the consumer holds the base pointer for the +/// tensor's whole lifetime. +pub fn export_to_dlpack_versioned(ndarray: NdArray) -> DLPackTensorVersioned { + let ndarray = if ndarray.data.is_owned() || ndarray.data.is_shared() { + ndarray + } else { + NdArray::from_buffer(ndarray.data.to_owned_copy(), ndarray.shape(), ndarray.strides()) + }; + let read_only = Arc::strong_count(&ndarray.data) > 1 || ndarray.data.is_shared(); + let flags = if read_only { DLPACK_FLAG_BITMASK_READ_ONLY } else { 0 }; + let shape = ndarray.shape().to_vec(); + let strides = ndarray.strides().to_vec(); + let data_ptr = ndarray.as_slice().as_ptr() as *mut c_void; + DLPackTensorVersioned::from_parts( + ndarray, data_ptr, 0, DLDataType::float::(), &shape, &strides, flags, + ) +} + +/// Export an NdArrayV window as a legacy DLPack managed tensor. +/// +/// Legacy `DLManagedTensor` carries no read-only flag, so a consumer +/// treats the pointer as writable. A view over a uniquely owned Minarrow +/// allocation exports zero-copy, carrying its offset and strides. A view +/// over shared or aliased storage materialises first so a legacy consumer +/// cannot mutate another reference. Consumers on the 1.x protocol should +/// prefer `to_dlpack_versioned`, which can share the backing and flag it +/// read-only. +#[cfg(feature = "views")] +pub fn export_view_to_dlpack(view: NdArrayV) -> DLPackTensor { + if !view.source.data.is_owned() || Arc::strong_count(&view.source.data) != 1 { + return export_to_dlpack(view.to_ndarray()); + } + let shape = view.shape().to_vec(); + let strides = view.strides().to_vec(); + let byte_offset = (view.offset * std::mem::size_of::()) as u64; + let data_ptr = view.source.as_slice().as_ptr() as *mut c_void; + DLPackTensor::from_parts(view, data_ptr, byte_offset, DLDataType::float::(), &shape, &strides) +} + +/// Export an NdArrayV window as a DLPack 1.x versioned managed tensor +/// without copying. The window offset carries through DLPack's +/// `byte_offset` field, and the view strides carry as element strides, +/// so sliced, transposed, and permuted views all hand over zero-copy. +/// A view shares its source buffer, so the read-only flag is set +/// whenever another reference to that buffer exists. A backing without +/// a stable owned or shared allocation materialises the window first. +#[cfg(feature = "views")] +pub fn export_view_to_dlpack_versioned(view: NdArrayV) -> DLPackTensorVersioned { + if !(view.source.data.is_owned() || view.source.data.is_shared()) { + return export_to_dlpack_versioned(view.to_ndarray()); + } + let read_only = Arc::strong_count(&view.source.data) > 1 || view.source.data.is_shared(); + let flags = if read_only { DLPACK_FLAG_BITMASK_READ_ONLY } else { 0 }; + let shape = view.shape().to_vec(); + let strides = view.strides().to_vec(); + let byte_offset = (view.offset * std::mem::size_of::()) as u64; + let data_ptr = view.source.as_slice().as_ptr() as *mut c_void; + DLPackTensorVersioned::from_parts( + view, data_ptr, byte_offset, DLDataType::float::(), &shape, &strides, flags, + ) +} + +/// Release callback invoked by the foreign consumer when it is done +/// with the tensor. Drops the holder which releases the source. +/// +/// # Safety +/// Must only be called once per DLManagedTensor. +unsafe extern "C" fn dlpack_deleter(managed: *mut DLManagedTensor) { + if managed.is_null() { return; } + let managed = unsafe { Box::from_raw(managed) }; + if !managed.manager_ctx.is_null() { + let _holder: Box> = unsafe { + Box::from_raw(managed.manager_ctx as *mut DLPackHolder) + }; + // holder drops here, releasing the source + } +} + +/// Release callback for the versioned managed tensor. +/// +/// # Safety +/// Must only be called once per DLManagedTensorVersioned. +unsafe extern "C" fn dlpack_deleter_versioned(managed: *mut DLManagedTensorVersioned) { + if managed.is_null() { return; } + let managed = unsafe { Box::from_raw(managed) }; + if !managed.manager_ctx.is_null() { + let _holder: Box> = unsafe { + Box::from_raw(managed.manager_ctx as *mut DLPackHolder) + }; + // holder drops here, releasing the source + } +} + +// **************************************************************** +// DLPackTensor / DLPackTensorVersioned - safe wrappers with Drop +// **************************************************************** + +/// Safe wrapper owning a legacy DLPack managed tensor. Calls the deleter +/// on drop. +/// +/// This is the return type of `NdArray::to_dlpack()`, analogous to how +/// `to_apache_arrow()` returns an `ArrayRef` that owns its lifecycle. +pub struct DLPackTensor { + ptr: *mut DLManagedTensor, +} + +impl DLPackTensor { + /// Assemble the managed tensor around a keep-alive source. + fn from_parts( + source: A, + data: *mut c_void, + byte_offset: u64, + dtype: DLDataType, + shape: &[usize], + strides: &[usize], + ) -> Self { + let mut holder = Box::new(DLPackHolder { + _source: source, + shape_i64: shape.iter().map(|&s| s as i64).collect(), + strides_i64: strides.iter().map(|&s| s as i64).collect(), + }); + let shape_ptr = if holder.shape_i64.is_empty() { + std::ptr::null_mut() + } else { + holder.shape_i64.as_mut_ptr() + }; + let strides_ptr = if holder.strides_i64.is_empty() { + std::ptr::null_mut() + } else { + holder.strides_i64.as_mut_ptr() + }; + let tensor = DLTensor { + data, + device: DLDevice::cpu(), + ndim: shape.len() as i32, + dtype, + shape: shape_ptr, + strides: strides_ptr, + byte_offset, + }; + let managed = Box::new(DLManagedTensor { + dl_tensor: tensor, + manager_ctx: Box::into_raw(holder) as *mut c_void, + deleter: Some(dlpack_deleter::), + }); + DLPackTensor { ptr: Box::into_raw(managed) } + } + + /// Raw pointer access for FFI consumers that need to take ownership + /// e.g. when creating a PyCapsule. + /// + /// After calling this, the DLPackTensor no longer owns the pointer + /// and will not call the deleter on drop. The caller is responsible for + /// ensuring the deleter is called. + pub fn into_raw(mut self) -> *mut DLManagedTensor { + let ptr = self.ptr; + self.ptr = std::ptr::null_mut(); + ptr + } + + /// Access the underlying DLTensor for reading shape, strides, data. + pub fn tensor(&self) -> &DLTensor { + unsafe { &(*self.ptr).dl_tensor } + } +} + +impl Drop for DLPackTensor { + fn drop(&mut self) { + if !self.ptr.is_null() { + unsafe { + if let Some(deleter) = (*self.ptr).deleter { + deleter(self.ptr); + } + } + } + } +} + +/// Safe wrapper owning a DLPack 1.x versioned managed tensor. Calls the +/// deleter on drop. +/// +/// This is the return type of `NdArray::to_dlpack_versioned()`, and the +/// payload a Python `__dlpack__` implementation places in a +/// `dltensor_versioned` capsule via `.into_raw()`. +pub struct DLPackTensorVersioned { + ptr: *mut DLManagedTensorVersioned, +} + +impl DLPackTensorVersioned { + /// Assemble the versioned managed tensor around a keep-alive source. + fn from_parts( + source: A, + data: *mut c_void, + byte_offset: u64, + dtype: DLDataType, + shape: &[usize], + strides: &[usize], + flags: u64, + ) -> Self { + let mut holder = Box::new(DLPackHolder { + _source: source, + shape_i64: shape.iter().map(|&s| s as i64).collect(), + strides_i64: strides.iter().map(|&s| s as i64).collect(), + }); + let shape_ptr = if holder.shape_i64.is_empty() { + std::ptr::null_mut() + } else { + holder.shape_i64.as_mut_ptr() + }; + let strides_ptr = if holder.strides_i64.is_empty() { + std::ptr::null_mut() + } else { + holder.strides_i64.as_mut_ptr() + }; + let tensor = DLTensor { + data, + device: DLDevice::cpu(), + ndim: shape.len() as i32, + dtype, + shape: shape_ptr, + strides: strides_ptr, + byte_offset, + }; + let managed = Box::new(DLManagedTensorVersioned { + version: DLPackVersion { + major: DLPACK_MAJOR_VERSION, + minor: DLPACK_MINOR_VERSION, + }, + manager_ctx: Box::into_raw(holder) as *mut c_void, + deleter: Some(dlpack_deleter_versioned::), + flags, + dl_tensor: tensor, + }); + DLPackTensorVersioned { ptr: Box::into_raw(managed) } + } + + /// Raw pointer access for FFI consumers that need to take ownership + /// e.g. when creating a PyCapsule. + /// + /// After calling this, the DLPackTensorVersioned no longer owns the + /// pointer and will not call the deleter on drop. The caller is + /// responsible for ensuring the deleter is called. + pub fn into_raw(mut self) -> *mut DLManagedTensorVersioned { + let ptr = self.ptr; + self.ptr = std::ptr::null_mut(); + ptr + } + + /// Access the underlying DLTensor for reading shape, strides, data. + pub fn tensor(&self) -> &DLTensor { + unsafe { &(*self.ptr).dl_tensor } + } + + /// The DLPack spec version stamped on this tensor. + pub fn version(&self) -> DLPackVersion { + unsafe { (*self.ptr).version } + } + + /// The flags word, including `DLPACK_FLAG_BITMASK_READ_ONLY`. + pub fn flags(&self) -> u64 { + unsafe { (*self.ptr).flags } + } +} + +impl Drop for DLPackTensorVersioned { + fn drop(&mut self) { + if !self.ptr.is_null() { + unsafe { + if let Some(deleter) = (*self.ptr).deleter { + deleter(self.ptr); + } + } + } + } +} + +// **************************************************************** +// Import: managed tensor -> NdArray +// **************************************************************** + +/// The foreign managed tensor behind an import, either ABI. +enum ForeignManaged { + Legacy(*mut DLManagedTensor), + Versioned(*mut DLManagedTensorVersioned), +} + +/// Owns the foreign managed tensor and calls its deleter on drop. +/// +/// The DLPack contract makes the deleter responsible for releasing the +/// managed tensor allocation, so the raw pointer is held rather than a +/// `Box`. Wrapping it in a `Box` would free the allocation a second time +/// after the deleter already released it. +struct ForeignDLPack { + ptr: *const u8, + len_bytes: usize, + managed: ForeignManaged, +} + +impl AsRef<[u8]> for ForeignDLPack { + fn as_ref(&self) -> &[u8] { + if self.len_bytes == 0 || self.ptr.is_null() { + return &[]; + } + unsafe { std::slice::from_raw_parts(self.ptr, self.len_bytes) } + } +} + +impl Drop for ForeignDLPack { + fn drop(&mut self) { + unsafe { + match self.managed { + ForeignManaged::Legacy(m) if !m.is_null() => { + if let Some(deleter) = (*m).deleter { + deleter(m); + } + } + ForeignManaged::Versioned(m) if !m.is_null() => { + if let Some(deleter) = (*m).deleter { + deleter(m); + } + } + _ => {} + } + } + self.managed = ForeignManaged::Legacy(std::ptr::null_mut()); + } +} + +unsafe impl Send for ForeignDLPack {} +unsafe impl Sync for ForeignDLPack {} + +/// Validation and wrapping pipeline shared by both import entry points. +/// +/// The guard owns the foreign tensor, so any validation failure drops it +/// and calls the foreign deleter as the ownership-transfer contract +/// requires. On the success path the guard moves into the SharedBuffer +/// and the deleter runs when the NdArray is dropped. +/// +/// # Safety +/// The DLTensor must belong to the managed tensor held by `foreign` and +/// point to accessible CPU memory. +unsafe fn import_dl_tensor( + tensor: &DLTensor, + mut foreign: ForeignDLPack, +) -> Result, MinarrowError> { + if tensor.device.device_type != DLDeviceType::CPU { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!( + "only CPU tensors are supported, got device type {}", + tensor.device.device_type.0 + ), + }); + } + + let bits = (std::mem::size_of::() * 8) as u8; + if tensor.dtype.code != 2 || tensor.dtype.bits != bits || tensor.dtype.lanes != 1 { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!( + "dtype mismatch, expected code=2 bits={}, got code={} bits={} lanes={}", + bits, tensor.dtype.code, tensor.dtype.bits, tensor.dtype.lanes + ), + }); + } + + if tensor.ndim < 0 { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!("tensor rank cannot be negative, got ndim={}", tensor.ndim), + }); + } + let ndim = tensor.ndim as usize; + + if ndim > 0 && tensor.shape.is_null() { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!("null shape pointer for ndim={}", ndim), + }); + } + let raw_shape: &[i64] = if ndim == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(tensor.shape, ndim) } + }; + if raw_shape.iter().any(|&s| s < 0) { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!("negative dimension in shape {:?}", raw_shape), + }); + } + let shape: Vec = raw_shape.iter().map(|&s| s as usize).collect(); + + // Null strides mean C-contiguous i.e. row-major. + let strides: Vec = if tensor.strides.is_null() { + let mut s = vec![1usize; ndim]; + if ndim > 0 { + for d in (0..ndim - 1).rev() { + s[d] = s[d + 1] * shape[d + 1]; + } + } + s + } else { + let raw_strides = unsafe { std::slice::from_raw_parts(tensor.strides, ndim) }; + if raw_strides.iter().any(|&s| s < 0) { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!( + "negative strides {:?} are not representable in NdArray", + raw_strides + ), + }); + } + raw_strides.iter().map(|&s| s as usize).collect() + }; + + // Buffer length required to reach the last element, overflow-checked. + let buf_len = if shape.iter().any(|&s| s == 0) { + 0 + } else { + let mut max_offset = 0usize; + for (&s, &st) in shape.iter().zip(strides.iter()) { + let contribution = (s - 1).checked_mul(st).ok_or_else(|| MinarrowError::BridgeError { + source: "dlpack", + message: format!("shape {:?} with strides {:?} overflows the addressable range", shape, strides), + })?; + max_offset = max_offset.checked_add(contribution).ok_or_else(|| MinarrowError::BridgeError { + source: "dlpack", + message: format!("shape {:?} with strides {:?} overflows the addressable range", shape, strides), + })?; + } + max_offset + 1 + }; + + // The byte window the elements span, overflow-checked and capped at + // isize::MAX per the slice safety contract. + let len_bytes = buf_len + .checked_mul(std::mem::size_of::()) + .filter(|&n| n <= isize::MAX as usize) + .ok_or_else(|| MinarrowError::BridgeError { + source: "dlpack", + message: format!( + "shape {:?} with strides {:?} spans more than isize::MAX bytes", + shape, strides + ), + })?; + + if tensor.data.is_null() && buf_len > 0 { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!("null data pointer for a tensor of {} elements", buf_len), + }); + } + + // Data pointer with byte offset, checked for element alignment. The + // 64-byte SIMD alignment is handled downstream - Buffer::from_shared + // copies a non-aligned foreign buffer into an aligned Vec64. + let data_ptr = if tensor.data.is_null() { + std::ptr::null() + } else { + unsafe { (tensor.data as *const u8).add(tensor.byte_offset as usize) } + }; + if buf_len > 0 && (data_ptr as usize) % std::mem::align_of::() != 0 { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!("data pointer is not aligned for a {}-bit element", bits), + }); + } + + // Record the data window now that shape and strides are validated. + foreign.ptr = data_ptr; + foreign.len_bytes = len_bytes; + + // Wrap into SharedBuffer -> Buffer -> NdArray. The guard moves into the + // SharedBuffer, so the deleter now runs when the NdArray is dropped. + let shared = SharedBuffer::from_owner(foreign); + let buffer: Buffer = Buffer::from_shared(shared); + + Ok(NdArray::from_buffer(buffer, &shape, &strides)) +} + +/// Import a foreign legacy DLManagedTensor as a Minarrow NdArray. +/// +/// Only CPU float tensors matching the element type `T` are supported. The +/// NdArray takes ownership of the DLManagedTensor and will call its deleter +/// when dropped. +/// +/// # Safety +/// - The DLManagedTensor must be valid and point to accessible CPU memory +/// - The caller transfers ownership - the tensor must not be used after this call +/// - The tensor must contain data of element type `T` +/// - The imported NdArray is Send + Sync, so the tensor's deleter must +/// tolerate running on any thread +pub unsafe fn import_from_dlpack( + managed_ptr: *mut DLManagedTensor, +) -> Result, MinarrowError> { + if managed_ptr.is_null() { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: "null DLManagedTensor pointer".to_string(), + }); + } + let foreign = ForeignDLPack { + ptr: std::ptr::null(), + len_bytes: 0, + managed: ForeignManaged::Legacy(managed_ptr), + }; + let tensor = unsafe { &(*managed_ptr).dl_tensor }; + unsafe { import_dl_tensor(tensor, foreign) } +} + +/// Import a foreign DLPack 1.x versioned managed tensor as a Minarrow +/// NdArray. +/// +/// Rejects tensors stamped with a major version newer than +/// [`DLPACK_MAJOR_VERSION`]. A tensor carrying the read-only flag lands +/// as a shared buffer, so the first mutation copies rather than writing +/// through to the producer's memory. +/// +/// # Safety +/// - The DLManagedTensorVersioned must be valid and point to accessible CPU memory +/// - The caller transfers ownership - the tensor must not be used after this call +/// - The tensor must contain data of element type `T` +/// - The imported NdArray is Send + Sync, so the tensor's deleter must +/// tolerate running on any thread +pub unsafe fn import_from_dlpack_versioned( + managed_ptr: *mut DLManagedTensorVersioned, +) -> Result, MinarrowError> { + if managed_ptr.is_null() { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: "null DLManagedTensorVersioned pointer".to_string(), + }); + } + let foreign = ForeignDLPack { + ptr: std::ptr::null(), + len_bytes: 0, + managed: ForeignManaged::Versioned(managed_ptr), + }; + let version = unsafe { (*managed_ptr).version }; + if version.major > DLPACK_MAJOR_VERSION { + return Err(MinarrowError::BridgeError { + source: "dlpack", + message: format!( + "DLPack major version {} is newer than the supported {}", + version.major, DLPACK_MAJOR_VERSION + ), + }); + } + let tensor = unsafe { &(*managed_ptr).dl_tensor }; + unsafe { import_dl_tensor(tensor, foreign) } +} + +// **************************************************************** +// Tests +// **************************************************************** + +#[cfg(test)] +mod tests { + use super::*; + #[cfg(all(feature = "views", feature = "select"))] + use crate::traits::selection::RowSelection; + + #[test] + fn export_import_rank_zero_scalar() { + let arr = NdArray::from_slice(&[42.0], &[]); + let owned = export_to_dlpack(arr); + let tensor = owned.tensor(); + assert_eq!(tensor.ndim, 0); + assert!(tensor.shape.is_null()); + assert!(tensor.strides.is_null()); + assert_eq!(unsafe { *(tensor.data as *const f64) }, 42.0); + + let imported = unsafe { import_from_dlpack::(owned.into_raw()) }.unwrap(); + assert!(imported.shape().is_empty()); + assert!(imported.strides().is_empty()); + assert_eq!(imported.len(), 1); + assert_eq!(imported.get(&[]), 42.0); + + let versioned = export_to_dlpack_versioned(NdArray::from_slice(&[7.0f32], &[])); + assert_eq!(versioned.tensor().ndim, 0); + assert!(versioned.tensor().shape.is_null()); + assert!(versioned.tensor().strides.is_null()); + let imported = unsafe { + import_from_dlpack_versioned::(versioned.into_raw()) + }.unwrap(); + assert!(imported.shape().is_empty()); + assert_eq!(imported.get(&[]), 7.0); + } + + #[test] + fn export_roundtrip_1d() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let owned = export_to_dlpack(arr); + let tensor = owned.tensor(); + + assert_eq!(tensor.ndim, 1); + assert_eq!(tensor.dtype.code, 2); + assert_eq!(tensor.dtype.bits, 64); + assert_eq!(tensor.byte_offset, 0); + + let shape = unsafe { std::slice::from_raw_parts(tensor.shape, 1) }; + assert_eq!(shape[0], 3); + + let strides = unsafe { std::slice::from_raw_parts(tensor.strides, 1) }; + assert_eq!(strides[0], 1); + + let data = tensor.data as *const f64; + assert_eq!(unsafe { *data }, 1.0); + assert_eq!(unsafe { *data.add(1) }, 2.0); + assert_eq!(unsafe { *data.add(2) }, 3.0); + // owned drops here, calling deleter + } + + #[test] + fn export_2d_column_major_strides() { + let arr = NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2] + ); + let owned = export_to_dlpack(arr); + let tensor = owned.tensor(); + + assert_eq!(tensor.ndim, 2); + + let shape = unsafe { std::slice::from_raw_parts(tensor.shape, 2) }; + assert_eq!(shape[0], 3); + assert_eq!(shape[1], 2); + + let strides = unsafe { std::slice::from_raw_parts(tensor.strides, 2) }; + assert_eq!(strides[0], 1); + assert_eq!(strides[1], 3); + + let data = tensor.data as *const f64; + assert_eq!(unsafe { *data }, 1.0); + assert_eq!(unsafe { *data.add(2) }, 3.0); + assert_eq!(unsafe { *data.add(strides[1] as usize) }, 4.0); + assert_eq!(unsafe { *data.add(strides[1] as usize + 2) }, 6.0); + } + + #[test] + fn export_preserves_data_lifetime() { + let owned = { + let arr = NdArray::from_slice(&[42.0, 99.0], &[2]); + export_to_dlpack(arr) + }; + let data = owned.tensor().data as *const f64; + assert_eq!(unsafe { *data }, 42.0); + assert_eq!(unsafe { *data.add(1) }, 99.0); + } + + #[test] + fn import_from_exported() { + let original = NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2] + ); + let owned = export_to_dlpack(original); + let raw = owned.into_raw(); + + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.get(&[0, 0]), 1.0); + assert_eq!(imported.get(&[2, 0]), 3.0); + assert_eq!(imported.get(&[0, 1]), 4.0); + assert_eq!(imported.get(&[2, 1]), 6.0); + } + + #[test] + fn import_null_ptr_fails() { + let result = unsafe { import_from_dlpack::(std::ptr::null_mut()) }; + assert!(result.is_err()); + } + + #[test] + fn export_3d_strides() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let arr = NdArray::from_slice(&data, &[2, 3, 4]); + let owned = export_to_dlpack(arr); + let tensor = owned.tensor(); + + assert_eq!(tensor.ndim, 3); + + let shape = unsafe { std::slice::from_raw_parts(tensor.shape, 3) }; + assert_eq!(shape, &[2, 3, 4]); + + let strides = unsafe { std::slice::from_raw_parts(tensor.strides, 3) }; + assert_eq!(strides[0], 1); + assert_eq!(strides[1], 2); + assert_eq!(strides[2], 6); + } + + #[test] + fn to_dlpack_method() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let owned = arr.to_dlpack(); + let tensor = owned.tensor(); + assert_eq!(tensor.ndim, 2); + let data = tensor.data as *const f64; + assert_eq!(unsafe { *data }, 1.0); + } + + #[test] + fn export_roundtrip_f32() { + let original = NdArray::::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2] + ); + let owned = export_to_dlpack(original); + let raw = owned.into_raw(); + + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.get(&[0, 0]), 1.0); + assert_eq!(imported.get(&[2, 0]), 3.0); + assert_eq!(imported.get(&[0, 1]), 4.0); + assert_eq!(imported.get(&[2, 1]), 6.0); + } + + #[test] + fn export_f32_dtype_bits() { + let arr = NdArray::::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let owned = export_to_dlpack(arr); + let tensor = owned.tensor(); + assert_eq!(tensor.dtype.code, 2); + assert_eq!(tensor.dtype.bits, 32); + } + + #[test] + fn import_dtype_mismatch_fails() { + // Export f64, then import as f32 - the bit width mismatch rejects. + let f64_arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let raw = export_to_dlpack(f64_arr).into_raw(); + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + + // Export f32, then import as f64 - likewise rejected. + let f32_arr = NdArray::::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + let raw = export_to_dlpack(f32_arr).into_raw(); + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + // *** Versioned ABI *********************************************** + + #[test] + fn versioned_roundtrip() { + let original = NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2] + ); + let owned = export_to_dlpack_versioned(original); + assert_eq!(owned.version().major, DLPACK_MAJOR_VERSION); + assert_eq!(owned.version().minor, DLPACK_MINOR_VERSION); + assert_eq!(owned.flags(), 0); + + let raw = owned.into_raw(); + let imported = unsafe { import_from_dlpack_versioned::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.get(&[0, 0]), 1.0); + assert_eq!(imported.get(&[2, 1]), 6.0); + } + + #[test] + fn versioned_read_only_flag_on_shared_buffer() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let clone = arr.clone(); + let owned = export_to_dlpack_versioned(clone); + assert_ne!(owned.flags() & DLPACK_FLAG_BITMASK_READ_ONLY, 0); + // The original still reads its data after the export drops. + drop(owned); + assert_eq!(arr.get(&[0]), 1.0); + } + + #[test] + fn versioned_rejects_newer_major() { + let arr = NdArray::from_slice(&[1.0, 2.0], &[2]); + let raw = export_to_dlpack_versioned(arr).into_raw(); + unsafe { (*raw).version.major = DLPACK_MAJOR_VERSION + 1 }; + let result = unsafe { import_from_dlpack_versioned::(raw) }; + assert!(result.is_err()); + } + + // *** View export ************************************************* + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn uniquely_owned_view_export_carries_byte_offset() { + let (view, data_ptr) = { + let data: Vec = (1..=10).map(|x| x as f64).collect(); + let arr = NdArray::from_slice(&data, &[5, 2]); + let data_ptr = arr.as_slice().as_ptr(); + (arr.r(2..5), data_ptr) + }; + let owned = export_view_to_dlpack(view); + let tensor = owned.tensor(); + + assert_eq!(tensor.byte_offset, 2 * std::mem::size_of::() as u64); + assert_eq!(tensor.data as *const f64, data_ptr); + let shape = unsafe { std::slice::from_raw_parts(tensor.shape, 2) }; + assert_eq!(shape, &[3, 2]); + + let raw = owned.into_raw(); + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.get(&[0, 0]), 3.0); + assert_eq!(imported.get(&[2, 1]), 10.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn aliased_view_export_materialises() { + let data: Vec = (1..=10).map(|x| x as f64).collect(); + let arr = NdArray::from_slice(&data, &[5, 2]); + let owned = export_view_to_dlpack(arr.r(2..5)); + let tensor = owned.tensor(); + + assert_eq!(tensor.byte_offset, 0); + assert_ne!(tensor.data as *const f64, arr.as_slice().as_ptr()); + + let raw = owned.into_raw(); + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.get(&[0, 0]), arr.get(&[2, 0])); + assert_eq!(imported.get(&[2, 1]), arr.get(&[4, 1])); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn view_export_versioned_sets_read_only() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // The view shares the source buffer, so the export is read-only. + let owned = export_view_to_dlpack_versioned(arr.r(0..2)); + assert_ne!(owned.flags() & DLPACK_FLAG_BITMASK_READ_ONLY, 0); + } + + #[cfg(feature = "views")] + #[test] + fn transposed_view_export_roundtrip() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let raw = export_view_to_dlpack(arr.as_view().transpose()).into_raw(); + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[2, 3]); + assert_eq!(imported.get(&[0, 0]), arr.get(&[0, 0])); + assert_eq!(imported.get(&[1, 2]), arr.get(&[2, 1])); + } + + // *** Import validation ******************************************* + + #[test] + fn import_rejects_negative_shape() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let raw = export_to_dlpack(arr).into_raw(); + unsafe { *(*raw).dl_tensor.shape = -2 }; + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + #[test] + fn import_rejects_negative_strides() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let raw = export_to_dlpack(arr).into_raw(); + unsafe { *(*raw).dl_tensor.strides = -1 }; + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + #[test] + fn import_rejects_misaligned_byte_offset() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]); + let raw = export_to_dlpack(arr).into_raw(); + unsafe { (*raw).dl_tensor.byte_offset = 4 }; + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + #[test] + fn import_rejects_overflowing_extents() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let raw = export_to_dlpack(arr).into_raw(); + unsafe { + *(*raw).dl_tensor.shape = i64::MAX / 4; + *(*raw).dl_tensor.strides = i64::MAX / 4; + } + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + #[test] + fn import_rejects_unknown_device() { + let arr = NdArray::from_slice(&[1.0, 2.0], &[2]); + let raw = export_to_dlpack(arr).into_raw(); + // kDLCUDAManaged, a code outside the host-memory set. + unsafe { (*raw).dl_tensor.device.device_type = DLDeviceType::CUDA_MANAGED }; + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + } + + #[test] + fn import_null_strides_as_row_major() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let raw = export_to_dlpack(arr).into_raw(); + unsafe { (*raw).dl_tensor.strides = std::ptr::null_mut() }; + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.shape(), &[3, 2]); + assert_eq!(imported.strides(), &[2, 1]); + // Row-major reinterpretation of the buffer. + assert_eq!(imported.get(&[0, 0]), 1.0); + assert_eq!(imported.get(&[0, 1]), 2.0); + assert_eq!(imported.get(&[1, 0]), 3.0); + } + + // *** Foreign producer tensors ************************************ + + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Shape and strides storage plus the counter the deleter bumps, + /// keeping the foreign tensor's pointers alive until release. + struct ForeignCtx { + shape: Vec, + strides: Vec, + hits: &'static AtomicUsize, + } + + unsafe extern "C" fn counting_deleter(managed: *mut DLManagedTensor) { + if managed.is_null() { + return; + } + let managed = unsafe { Box::from_raw(managed) }; + let ctx: Box = + unsafe { Box::from_raw(managed.manager_ctx as *mut ForeignCtx) }; + ctx.hits.fetch_add(1, Ordering::SeqCst); + } + + /// Assemble a foreign-owned legacy managed tensor over caller-supplied + /// f64 data, wired to the counting deleter. + fn foreign_tensor( + data: *mut c_void, + shape: Vec, + strides: Vec, + hits: &'static AtomicUsize, + ) -> *mut DLManagedTensor { + let mut ctx = Box::new(ForeignCtx { shape, strides, hits }); + let tensor = DLTensor { + data, + device: DLDevice::cpu(), + ndim: ctx.shape.len() as i32, + dtype: DLDataType::FLOAT64, + shape: ctx.shape.as_mut_ptr(), + strides: ctx.strides.as_mut_ptr(), + byte_offset: 0, + }; + Box::into_raw(Box::new(DLManagedTensor { + dl_tensor: tensor, + manager_ctx: Box::into_raw(ctx) as *mut c_void, + deleter: Some(counting_deleter), + })) + } + + /// Allocate a 64-byte aligned f64 buffer the test controls outright, + /// so a zero-copy import shares it without a realignment copy. + fn aligned_f64_alloc(values: &[f64]) -> (*mut f64, std::alloc::Layout) { + let layout = std::alloc::Layout::from_size_align( + values.len() * std::mem::size_of::(), + 64, + ) + .unwrap(); + let ptr = unsafe { std::alloc::alloc(layout) as *mut f64 }; + assert!(!ptr.is_null()); + unsafe { std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len()) }; + (ptr, layout) + } + + #[test] + fn import_rejects_null_shape_pointer() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let raw = foreign_tensor(64 as *mut c_void, vec![2], vec![1], &HITS); + unsafe { (*raw).dl_tensor.shape = std::ptr::null_mut() }; + let result = unsafe { import_from_dlpack::(raw) }; + let Err(MinarrowError::BridgeError { message, .. }) = result else { + panic!("expected BridgeError"); + }; + assert!(message.contains("null shape pointer")); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + } + + #[test] + fn import_rejects_null_data_pointer() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let raw = foreign_tensor(std::ptr::null_mut(), vec![2], vec![1], &HITS); + let result = unsafe { import_from_dlpack::(raw) }; + let Err(MinarrowError::BridgeError { message, .. }) = result else { + panic!("expected BridgeError"); + }; + assert!(message.contains("null data pointer")); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + } + + #[test] + fn import_rejects_span_beyond_isize_max() { + static HITS: AtomicUsize = AtomicUsize::new(0); + // Validation rejects the span before any read, so the dangling + // but aligned data pointer is never dereferenced. + let raw = foreign_tensor(64 as *mut c_void, vec![i64::MAX / 4], vec![1], &HITS); + let result = unsafe { import_from_dlpack::(raw) }; + let Err(MinarrowError::BridgeError { message, .. }) = result else { + panic!("expected BridgeError"); + }; + assert!(message.contains("spans more than isize::MAX bytes")); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + } + + #[test] + fn deleter_runs_once_on_drop() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let (data, layout) = aligned_f64_alloc(&[1.0, 2.0, 3.0, 4.0]); + let raw = foreign_tensor(data as *mut c_void, vec![4], vec![1], &HITS); + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + assert_eq!(imported.get(&[2]), 3.0); + assert_eq!(HITS.load(Ordering::SeqCst), 0); + drop(imported); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + unsafe { std::alloc::dealloc(data as *mut u8, layout) }; + } + + #[test] + fn deleter_runs_once_on_failed_import() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let (data, layout) = aligned_f64_alloc(&[1.0, 2.0]); + let raw = foreign_tensor(data as *mut c_void, vec![2], vec![1], &HITS); + // Importing as f32 against the f64 payload fails the dtype check. + let result = unsafe { import_from_dlpack::(raw) }; + assert!(result.is_err()); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + unsafe { std::alloc::dealloc(data as *mut u8, layout) }; + } + + #[test] + fn import_mutation_copies_rather_than_writing_through() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let (data, layout) = aligned_f64_alloc(&[1.0, 2.0, 3.0, 4.0]); + let raw = foreign_tensor(data as *mut c_void, vec![4], vec![1], &HITS); + let mut imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + imported.set(&[1], 99.0); + assert_eq!(imported.get(&[1]), 99.0); + // The producer's memory holds its original values. + let original = unsafe { std::slice::from_raw_parts(data, 4) }; + assert_eq!(original, &[1.0, 2.0, 3.0, 4.0]); + drop(imported); + unsafe { std::alloc::dealloc(data as *mut u8, layout) }; + } + + #[test] + fn export_aliased_buffer_copies() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let alias = arr.clone(); + let exported = export_to_dlpack(arr); + assert_ne!( + exported.tensor().data as *const f64, + alias.as_slice().as_ptr() + ); + } + + #[test] + fn export_uniquely_owned_buffer_is_zero_copy() { + let arr = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let data_ptr = arr.as_slice().as_ptr(); + let exported = export_to_dlpack(arr); + assert_eq!(exported.tensor().data as *const f64, data_ptr); + } + + #[test] + fn imported_ndarray_send_drops_on_other_thread() { + static HITS: AtomicUsize = AtomicUsize::new(0); + let (data, layout) = aligned_f64_alloc(&[1.0, 2.0, 3.0, 4.0]); + let raw = foreign_tensor(data as *mut c_void, vec![4], vec![1], &HITS); + let imported = unsafe { import_from_dlpack::(raw) }.unwrap(); + std::thread::spawn(move || { + assert_eq!(imported.get(&[3]), 4.0); + drop(imported); + }) + .join() + .unwrap(); + assert_eq!(HITS.load(Ordering::SeqCst), 1); + unsafe { std::alloc::dealloc(data as *mut u8, layout) }; + } +} diff --git a/src/kernels/arithmetic/mod.rs b/src/kernels/arithmetic/mod.rs index a412103..62c3116 100644 --- a/src/kernels/arithmetic/mod.rs +++ b/src/kernels/arithmetic/mod.rs @@ -36,7 +36,7 @@ pub mod simd; pub mod std; pub mod string; pub mod string_ops; -#[cfg(feature = "broadcast")] +#[cfg(all(feature = "broadcast", feature = "value_type", feature = "views"))] pub mod types; // Shared tests for SIMD and Std diff --git a/src/kernels/arithmetic/types.rs b/src/kernels/arithmetic/types.rs index 12707d8..1895674 100644 --- a/src/kernels/arithmetic/types.rs +++ b/src/kernels/arithmetic/types.rs @@ -836,6 +836,7 @@ impl Rem for SuperTableV { } } + // Cube implementations #[cfg(feature = "cube")] impl Add for Cube { @@ -918,6 +919,67 @@ mod tests { use crate::traits::selection::ColumnSelection; use crate::{Array, IntegerArray, NumericArray, vec64}; + #[cfg(feature = "ndarray")] + #[test] + fn ndarray_family_values_are_not_arithmetic_operands() { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + use crate::NdArray; + + fn assert_unimplemented(f: impl FnOnce()) { + let panic = catch_unwind(AssertUnwindSafe(f)); + assert!(panic.is_err()); + } + + let a = NdArray::from_slice(&[1.0, 2.0], &[2]); + let lhs = Value::NdArray(Arc::new(a.clone())); + let rhs = Value::NdArray(Arc::new(a.clone())); + assert_unimplemented(|| { + let _ = lhs + rhs; + }); + + #[cfg(feature = "views")] + { + let lhs = Value::NdArrayView(Arc::new(a.as_view())); + let rhs = lhs.clone(); + assert_unimplemented(|| { + let _ = lhs + rhs; + }); + } + + #[cfg(feature = "chunked")] + { + use crate::SuperNdArray; + + let a = SuperNdArray::from_batches(vec![a.clone()], "values"); + let lhs = Value::SuperNdArray(Arc::new(a.clone())); + let rhs = lhs.clone(); + assert_unimplemented(|| { + let _ = lhs + rhs; + }); + + #[cfg(feature = "views")] + { + let lhs = Value::SuperNdArrayView(Arc::new(a.slice(0, 2))); + let rhs = lhs.clone(); + assert_unimplemented(|| { + let _ = lhs + rhs; + }); + } + } + + #[cfg(feature = "xarray")] + { + use crate::XArray; + + let a = Value::XArray(Arc::new(XArray::new(a, &["value"]))); + let b = a.clone(); + assert_unimplemented(|| { + let _ = a + b; + }); + } + } + #[test] fn test_value_addition() { let arr1 = diff --git a/src/kernels/broadcast/array.rs b/src/kernels/broadcast/array.rs index 2779dc8..3e6224a 100644 --- a/src/kernels/broadcast/array.rs +++ b/src/kernels/broadcast/array.rs @@ -26,7 +26,9 @@ use crate::kernels::routing::arithmetic::resolve_binary_arithmetic; use crate::structs::field_array::create_field_for_array; use crate::{Array, ArrayV, Bitmask, FieldArray, Table, Value}; #[cfg(feature = "scalar_type")] -use crate::{BooleanArray, DatetimeArray, FloatArray, IntegerArray, Scalar, StringArray}; +use crate::{BooleanArray, FloatArray, IntegerArray, Scalar, StringArray}; +#[cfg(all(feature = "datetime", feature = "scalar_type"))] +use crate::DatetimeArray; use std::sync::Arc; /// Broadcast addition: `lhs + rhs` with automatic scalar expansion. diff --git a/src/kernels/broadcast/array_view.rs b/src/kernels/broadcast/array_view.rs index 9f28ba4..4d013df 100644 --- a/src/kernels/broadcast/array_view.rs +++ b/src/kernels/broadcast/array_view.rs @@ -124,7 +124,7 @@ pub fn broadcast_arrayview_to_tableview( } /// Helper function for ArrayView-SuperTableView broadcasting - work per chunk by slicing the existing ArrayView -#[cfg(feature = "views")] +#[cfg(all(feature = "chunked", feature = "views"))] pub fn broadcast_arrayview_to_supertableview( op: ArithmeticOperator, array_view: &ArrayV, diff --git a/src/kernels/broadcast/cube.rs b/src/kernels/broadcast/cube.rs index e8d00ea..aa34884 100644 --- a/src/kernels/broadcast/cube.rs +++ b/src/kernels/broadcast/cube.rs @@ -33,8 +33,10 @@ use crate::enums::{error::MinarrowError, operators::ArithmeticOperator}; #[cfg(feature = "cube")] use crate::kernels::broadcast::{ array::broadcast_array_to_table, - table::{broadcast_table_to_array, broadcast_table_to_scalar, broadcast_table_with_operator}, + table::{broadcast_table_to_array, broadcast_table_with_operator}, }; +#[cfg(all(feature = "cube", feature = "scalar_type"))] +use crate::kernels::broadcast::table::broadcast_table_to_scalar; /// Broadcasts addition over Cube tables (3D structure) /// diff --git a/src/kernels/broadcast/mod.rs b/src/kernels/broadcast/mod.rs index 71bea93..2ce9122 100644 --- a/src/kernels/broadcast/mod.rs +++ b/src/kernels/broadcast/mod.rs @@ -22,6 +22,8 @@ //! //! This enables ergonomic arithmetic operations like: //! ```rust +//! # #[cfg(all(feature = "value_type", feature = "views"))] +//! # { //! use minarrow::{Value, arr_i32, vec64}; //! use std::sync::Arc; //! let arr1 = arr_i32![1, 2, 3, 4]; @@ -29,115 +31,230 @@ //! let a = Value::Array(Arc::new(arr1)); //! let b = Value::Array(Arc::new(arr2)); //! let result = a + b; // Automatically broadcasts and performs element-wise addition +//! # } //! ``` //! +#[cfg(all(feature = "value_type", feature = "views"))] pub mod array; +#[cfg(all(feature = "value_type", feature = "views"))] pub mod array_view; -#[cfg(feature = "cube")] +#[cfg(all(feature = "cube", feature = "value_type", feature = "views"))] pub mod cube; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] pub mod field_array; #[cfg(feature = "matrix")] pub mod matrix; +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] pub mod scalar; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] pub mod super_array; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] pub mod super_array_view; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] pub mod super_table; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] pub mod super_table_view; +#[cfg(all(feature = "value_type", feature = "views"))] pub mod table; +#[cfg(all(feature = "value_type", feature = "views"))] pub mod table_view; -#[cfg(feature = "chunked")] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use crate::utils::create_aligned_chunks_from_array; -pub use table::{broadcast_super_table_add, broadcast_table_add}; +#[cfg(all(feature = "value_type", feature = "views"))] +pub use table::broadcast_table_add; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] +pub use table::broadcast_super_table_add; // Import helper functions from submodules -#[cfg(feature = "scalar_type")] +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use crate::kernels::routing::arithmetic::scalar_arithmetic; -#[cfg(all(feature = "chunked", feature = "views"))] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use array::broadcast_array_to_supertableview; +#[cfg(all(feature = "value_type", feature = "views"))] use array::broadcast_array_to_table; -#[cfg(feature = "views")] -use array_view::{ - broadcast_arrayview_to_supertableview, broadcast_arrayview_to_table, - broadcast_arrayview_to_tableview, -}; -#[cfg(all(feature = "scalar_type", feature = "chunked", feature = "views"))] +#[cfg(all(feature = "value_type", feature = "views"))] +use array_view::{broadcast_arrayview_to_table, broadcast_arrayview_to_tableview}; +#[cfg(all(feature = "chunked", feature = "value_type", feature = "views"))] +use array_view::broadcast_arrayview_to_supertableview; +#[cfg(all( + feature = "chunked", + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use scalar::broadcast_scalar_to_supertableview; -#[cfg(feature = "scalar_type")] +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use scalar::broadcast_scalar_to_table; -#[cfg(all(feature = "scalar_type", feature = "views"))] +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use scalar::broadcast_scalar_to_tableview; -#[cfg(feature = "chunked")] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use super_array::broadcast_superarray_to_table; +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use super_array::route_super_array_broadcast; -#[cfg(all(feature = "chunked", feature = "views"))] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use super_array_view::broadcast_superarrayview_to_tableview; -#[cfg(feature = "chunked")] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use super_table::broadcast_super_table_with_operator; -#[cfg(all(feature = "scalar_type", feature = "chunked", feature = "views"))] +#[cfg(all( + feature = "chunked", + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use super_table_view::broadcast_supertableview_to_scalar; -#[cfg(all(feature = "chunked", feature = "views"))] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use super_table_view::{ broadcast_superarrayview_to_table, broadcast_supertableview_to_array, broadcast_supertableview_to_arrayview, }; -#[cfg(feature = "views")] +#[cfg(all(feature = "value_type", feature = "views"))] use table::broadcast_table_to_arrayview; -use table::{broadcast_table_to_array, broadcast_table_to_scalar, broadcast_table_with_operator}; -#[cfg(feature = "chunked")] +#[cfg(all(feature = "value_type", feature = "views"))] +use table::{broadcast_table_to_array, broadcast_table_with_operator}; +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] +use table::broadcast_table_to_scalar; +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use table::{broadcast_table_to_superarray, broadcast_table_to_superarrayview}; -#[cfg(all(feature = "scalar_type", feature = "views"))] +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use table_view::broadcast_tableview_to_scalar; -#[cfg(all(feature = "chunked", feature = "views"))] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use table_view::broadcast_tableview_to_superarrayview; -#[cfg(feature = "views")] +#[cfg(all(feature = "value_type", feature = "views"))] use table_view::{broadcast_tableview_to_arrayview, broadcast_tableview_to_tableview}; -#[cfg(feature = "cube")] +#[cfg(all( + feature = "cube", + feature = "value_type", + feature = "views" +))] use crate::Cube; -#[cfg(feature = "views")] +#[cfg(all(feature = "value_type", feature = "views"))] use crate::ArrayV; +#[cfg(all(feature = "value_type", feature = "views"))] use crate::{Array, FloatArray, IntegerArray, StringArray}; -#[cfg(feature = "scalar_type")] +#[cfg(all( + feature = "scalar_type", + feature = "value_type", + feature = "views" +))] use crate::Scalar; -#[cfg(all(feature = "views", feature = "chunked"))] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use crate::SuperTableV; +#[cfg(all(feature = "value_type", feature = "views"))] use crate::enums::error::MinarrowError; +#[cfg(all(feature = "value_type", feature = "views"))] use crate::enums::operators::ArithmeticOperator; +#[cfg(all(feature = "value_type", feature = "views"))] use crate::enums::value::Value; -#[cfg(feature = "chunked")] +#[cfg(all( + feature = "chunked", + feature = "value_type", + feature = "views" +))] use crate::{SuperArray, SuperTable}; +#[cfg(all(feature = "value_type", feature = "views"))] use crate::kernels::routing::arithmetic::resolve_binary_arithmetic; /// Add two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_add(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Add, lhs, rhs) } /// Subtract two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_subtract(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Subtract, lhs, rhs) } /// Multiply two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_multiply(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Multiply, lhs, rhs) } /// Divide two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_divide(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Divide, lhs, rhs) } /// Remainder (modulo) two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_remainder(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Remainder, lhs, rhs) } /// Power/exponentiation of two Values with automatic broadcasting +#[cfg(all(feature = "value_type", feature = "views"))] pub fn value_power(lhs: Value, rhs: Value) -> Result { broadcast_value(ArithmeticOperator::Power, lhs, rhs) } @@ -148,7 +265,7 @@ pub fn value_power(lhs: Value, rhs: Value) -> Result { /// 1.⚠️ Best to keep this out of the binary by disabling value_type unless you /// require universal broadcasting compatibility. /// 2.These do not yet implement parallel processing to speed up broadcasting. -#[cfg(all(feature = "scalar_type", feature = "value_type"))] +#[cfg(all(feature = "value_type", feature = "views"))] pub fn broadcast_value( op: ArithmeticOperator, lhs: Value, @@ -410,6 +527,34 @@ pub fn broadcast_value( ), }), + // NdArray is not implemented at the present time. + #[cfg(feature = "ndarray")] + (Value::NdArray(_), _) | (_, Value::NdArray(_)) => { + unimplemented!("NdArray broadcasting") + } + + #[cfg(all(feature = "ndarray", feature = "views"))] + (Value::NdArrayView(_), _) | (_, Value::NdArrayView(_)) => { + unimplemented!("NdArrayView broadcasting") + } + + #[cfg(all(feature = "ndarray", feature = "chunked"))] + (Value::SuperNdArray(_), _) | (_, Value::SuperNdArray(_)) => { + unimplemented!("SuperNdArray broadcasting") + } + + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + (Value::SuperNdArrayView(_), _) | (_, Value::SuperNdArrayView(_)) => { + unimplemented!("SuperNdArrayView broadcasting") + } + + + // XArray is not implemented at the present time. + #[cfg(feature = "xarray")] + (Value::XArray(_), _) | (_, Value::XArray(_)) => { + unimplemented!("XArray broadcasting") + } + // Cube cases - use broadcast_cube_add #[cfg(feature = "cube")] (Value::Cube(l), Value::Cube(r)) => { @@ -497,11 +642,11 @@ pub fn broadcast_value( }, // Scalar-SuperTable broadcasting - apply to each table in SuperTable - #[cfg(feature = "scalar_type")] + #[cfg(all(feature = "chunked", feature = "scalar_type"))] (Value::Scalar(scalar), Value::SuperTable(super_table)) => { scalar::broadcast_scalar_to_supertable(op, &scalar, &super_table).map(|st| Value::SuperTable(Arc::new(st))) }, - #[cfg(feature = "scalar_type")] + #[cfg(all(feature = "chunked", feature = "scalar_type"))] (Value::SuperTable(super_table), Value::Scalar(scalar)) => { super_table::broadcast_supertable_to_scalar(op, &super_table, &scalar).map(|st| Value::SuperTable(Arc::new(st))) }, @@ -554,9 +699,11 @@ pub fn broadcast_value( }, // Array-SuperTable broadcasting - apply to each table in SuperTable + #[cfg(feature = "chunked")] (Value::Array(array), Value::SuperTable(super_table)) => { array::broadcast_array_to_supertable(op, &array, &super_table).map(|st| Value::SuperTable(Arc::new(st))) }, + #[cfg(feature = "chunked")] (Value::SuperTable(super_table), Value::Array(array)) => { super_table::broadcast_supertable_to_array(op, &super_table, &array).map(|st| Value::SuperTable(Arc::new(st))) }, @@ -570,9 +717,11 @@ pub fn broadcast_value( }, // FieldArray-SuperTable broadcasting - apply to each table in SuperTable + #[cfg(feature = "chunked")] (Value::FieldArray(field_array), Value::SuperTable(super_table)) => { super_table::broadcast_fieldarray_to_supertable(op, &field_array, &super_table).map(|st| Value::SuperTable(Arc::new(st))) }, + #[cfg(feature = "chunked")] (Value::SuperTable(super_table), Value::FieldArray(field_array)) => { super_table::broadcast_supertable_to_fieldarray(op, &super_table, &field_array).map(|st| Value::SuperTable(Arc::new(st))) }, @@ -588,11 +737,11 @@ pub fn broadcast_value( }, // ArrayView-SuperTable broadcasting - convert views to arrays and broadcast - #[cfg(feature = "views")] + #[cfg(all(feature = "chunked", feature = "views"))] (Value::ArrayView(array_view), Value::SuperTable(super_table)) => { super_table::broadcast_arrayview_to_supertable(op, &array_view, &super_table).map(|st| Value::SuperTable(Arc::new(st))) }, - #[cfg(feature = "views")] + #[cfg(all(feature = "chunked", feature = "views"))] (Value::SuperTable(super_table), Value::ArrayView(array_view)) => { super_table::broadcast_supertable_to_arrayview(op, &super_table, &array_view).map(|st| Value::SuperTable(Arc::new(st))) }, @@ -608,11 +757,11 @@ pub fn broadcast_value( }, // ArrayView-SuperTableView broadcasting - work per chunk, not materialised - #[cfg(feature = "views")] + #[cfg(all(feature = "chunked", feature = "views"))] (Value::ArrayView(array_view), Value::SuperTableView(super_table_view)) => { broadcast_arrayview_to_supertableview(op, &array_view, &super_table_view).map(|stv| Value::SuperTableView(Arc::new(stv))) }, - #[cfg(feature = "views")] + #[cfg(all(feature = "chunked", feature = "views"))] (Value::SuperTableView(super_table_view), Value::ArrayView(array_view)) => { broadcast_supertableview_to_arrayview(op, &super_table_view, &array_view).map(|stv| Value::SuperTableView(Arc::new(stv))) }, diff --git a/src/kernels/broadcast/scalar.rs b/src/kernels/broadcast/scalar.rs index d434f98..2ca0832 100644 --- a/src/kernels/broadcast/scalar.rs +++ b/src/kernels/broadcast/scalar.rs @@ -22,11 +22,15 @@ use crate::kernels::broadcast::broadcast_value; use crate::kernels::routing::arithmetic::resolve_binary_arithmetic; use crate::structs::field_array::create_field_for_array; use crate::{ - Array, BooleanArray, CategoricalArray, DatetimeArray, FieldArray, FloatArray, IntegerArray, - Scalar, StringArray, Table, TableV, TextArray, Value, + Array, BooleanArray, CategoricalArray, FieldArray, FloatArray, IntegerArray, Scalar, + StringArray, Table, TableV, TextArray, Value, }; +#[cfg(feature = "datetime")] +use crate::DatetimeArray; #[cfg(feature = "views")] -use crate::{NumericArrayV, TemporalArrayV, TextArrayV}; +use crate::{NumericArrayV, TextArrayV}; +#[cfg(all(feature = "datetime", feature = "views"))] +use crate::TemporalArrayV; #[cfg(feature = "chunked")] use crate::{SuperArray, SuperTable, SuperTableV}; use std::sync::Arc; diff --git a/src/kernels/broadcast/super_table.rs b/src/kernels/broadcast/super_table.rs index 25671c2..5563282 100644 --- a/src/kernels/broadcast/super_table.rs +++ b/src/kernels/broadcast/super_table.rs @@ -15,9 +15,9 @@ use crate::enums::error::MinarrowError; use crate::enums::operators::ArithmeticOperator; use crate::kernels::broadcast::array::broadcast_array_to_table; -use crate::kernels::broadcast::table::{ - broadcast_table_to_array, broadcast_table_to_scalar, broadcast_table_with_operator, -}; +use crate::kernels::broadcast::table::{broadcast_table_to_array, broadcast_table_with_operator}; +#[cfg(feature = "scalar_type")] +use crate::kernels::broadcast::table::broadcast_table_to_scalar; use crate::{SuperTable, SuperTableV, Value}; use std::sync::Arc; diff --git a/src/kernels/broadcast/super_table_view.rs b/src/kernels/broadcast/super_table_view.rs index 98d8cf4..30950a4 100644 --- a/src/kernels/broadcast/super_table_view.rs +++ b/src/kernels/broadcast/super_table_view.rs @@ -12,16 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "scalar_type")] use std::sync::Arc; use crate::enums::error::MinarrowError; use crate::enums::operators::ArithmeticOperator; use crate::kernels::broadcast::array_view::broadcast_arrayview_to_tableview; +#[cfg(feature = "scalar_type")] use crate::kernels::broadcast::broadcast_value; use crate::kernels::broadcast::table_view::{ broadcast_tableview_to_arrayview, broadcast_tableview_to_tableview, }; -use crate::{Array, ArrayV, Scalar, SuperArrayV, SuperTableV, Table, TableV, Value}; +use crate::{Array, ArrayV, SuperArrayV, SuperTableV, Table, TableV}; +#[cfg(feature = "scalar_type")] +use crate::{Scalar, Value}; /// Helper function for supertableview-scalar broadcasting - convert to table, broadcast, return as table #[cfg(all(feature = "scalar_type", feature = "chunked", feature = "views"))] diff --git a/src/kernels/broadcast/table_view.rs b/src/kernels/broadcast/table_view.rs index 305ab37..8972d09 100644 --- a/src/kernels/broadcast/table_view.rs +++ b/src/kernels/broadcast/table_view.rs @@ -18,7 +18,11 @@ use crate::enums::error::MinarrowError; use crate::enums::operators::ArithmeticOperator; use crate::kernels::broadcast::broadcast_value; use crate::kernels::routing::arithmetic::resolve_binary_arithmetic; -use crate::{ArrayV, FieldArray, Scalar, SuperArrayV, SuperTableV, Table, TableV, Value}; +use crate::{ArrayV, FieldArray, Table, TableV, Value}; +#[cfg(feature = "scalar_type")] +use crate::Scalar; +#[cfg(feature = "chunked")] +use crate::{SuperArrayV, SuperTableV}; /// Helper function for TableView-TableView broadcasting - work directly with views #[cfg(feature = "views")] diff --git a/src/lib.rs b/src/lib.rs index 01016bd..2a94c4f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -164,6 +164,8 @@ pub mod structs { #[cfg(feature = "chunked")] pub mod chunked { pub mod super_array; + #[cfg(feature = "ndarray")] + pub mod super_ndarray; pub mod super_table; } @@ -181,6 +183,8 @@ pub mod structs { #[cfg(feature = "chunked")] pub mod chunked { pub mod super_array_view; + #[cfg(feature = "ndarray")] + pub mod super_ndarray_view; pub mod super_table_view; } #[cfg(feature = "views")] @@ -194,6 +198,9 @@ pub mod structs { #[cfg(feature = "views")] pub mod array_view; pub mod bitmask_view; + #[cfg(feature = "views")] + #[cfg(feature = "ndarray")] + pub mod ndarray_view; #[cfg(feature = "views")] pub mod table_view; @@ -211,8 +218,12 @@ pub mod structs { pub mod lbuffer; #[cfg(feature = "matrix")] pub mod matrix; + #[cfg(feature = "ndarray")] + pub mod ndarray; pub mod shared_buffer; pub mod table; + #[cfg(feature = "xarray")] + pub mod xarray; } /// **Shared Memory** - *For sending data to other runtime(s) over FFI.* @@ -223,6 +234,8 @@ pub mod ffi { pub mod arrow_rs; #[cfg(feature = "cast_polars")] pub mod polars; + #[cfg(feature = "dlpack")] + pub mod dlpack; pub mod schema; } @@ -280,6 +293,8 @@ pub use structs::chunked::{ super_array::{RechunkStrategy, SuperArray}, super_table::SuperTable, }; +#[cfg(all(feature = "chunked", feature = "ndarray"))] +pub use structs::chunked::super_ndarray::SuperNdArray; #[cfg(feature = "lbuffer")] pub use structs::lbuffer::{LBuffer, LBufferV}; #[cfg(feature = "views")] @@ -288,6 +303,8 @@ pub use structs::views::bitmask_view::BitmaskV; #[cfg(feature = "views")] #[cfg(feature = "chunked")] pub use structs::views::chunked::{super_array_view::SuperArrayV, super_table_view::SuperTableV}; +#[cfg(all(feature = "views", feature = "chunked", feature = "ndarray"))] +pub use structs::views::chunked::super_ndarray_view::SuperNdArrayV; #[cfg(feature = "views")] pub use structs::views::collections::boolean_array_view::BooleanArrayV; #[cfg(feature = "views")] @@ -297,8 +314,12 @@ pub use structs::views::collections::numeric_array_view::NumericArrayV; pub use structs::views::collections::temporal_array_view::TemporalArrayV; #[cfg(feature = "views")] pub use structs::views::collections::text_array_view::TextArrayV; +#[cfg(all(feature = "views", feature = "ndarray"))] +pub use structs::views::ndarray_view::NdArrayV; pub use ffi::arrow_dtype::ArrowType; +#[cfg(feature = "dlpack")] +pub use ffi::dlpack::{DLPackTensor, DLPackTensorVersioned}; pub use structs::column::{Column, column}; #[cfg(feature = "cube")] pub use structs::cube::Cube; @@ -308,6 +329,10 @@ pub use structs::field::Field; pub use structs::field_array::{FieldArray, field_array}; #[cfg(feature = "matrix")] pub use structs::matrix::Matrix; +#[cfg(feature = "ndarray")] +pub use structs::ndarray::NdArray; +#[cfg(feature = "xarray")] +pub use structs::xarray::XArray; pub use structs::shared_buffer::SharedBuffer; pub use structs::table::Table; pub use structs::variants::boolean::BooleanArray; @@ -331,6 +356,8 @@ pub use traits::consolidate::Consolidate; pub use traits::datetime_ops::DatetimeOps; pub use traits::masked_array::MaskedArray; pub use traits::print::Print; +#[cfg(all(feature = "select", feature = "ndarray"))] +pub use traits::selection::AxisSelection; #[cfg(feature = "select")] pub use traits::selection::{ColumnSelection, RowSelection, Selection2D}; pub use traits::type_unions::{Float, Integer, Numeric, Primitive}; diff --git a/src/structs/buffer.rs b/src/structs/buffer.rs index dadbd8f..7a7d3cc 100644 --- a/src/structs/buffer.rs +++ b/src/structs/buffer.rs @@ -677,6 +677,14 @@ impl Buffer { matches!(self.storage, Storage::Shared { .. }) } + /// Returns true if the buffer owns its allocation as a `Vec64`. + /// Owned and shared storage hold a stable base pointer for the + /// allocation's lifetime, which FFI export relies on. + #[inline] + pub fn is_owned(&self) -> bool { + matches!(self.storage, Storage::Owned(_)) + } + /// Creates an owned copy of the buffer data. /// If the buffer is already owned, this clones the data. /// If the buffer is shared, this copies the data into a new owned Vec64. diff --git a/src/structs/chunked/super_ndarray.rs b/src/structs/chunked/super_ndarray.rs new file mode 100644 index 0000000..d2afd34 --- /dev/null +++ b/src/structs/chunked/super_ndarray.rs @@ -0,0 +1,1708 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! # **SuperNdArray** - *N-dimensional data in batches* +//! +//! The N-dimensional equivalent of [`SuperTable`](crate::SuperTable): it keeps +//! multiple `NdArray` batches under a shared rank and trailing shape. Batch +//! lengths may differ along axis 0, while indexing, iteration, and selection +//! use one combined axis-0 coordinate space. +//! +//! ## Typical uses +//! - Sensor or telemetry data arriving as independent temporal batches. +//! - Partitioned reads from storage where batches arrive independently. +//! - Memory-bounded ingestion where you want to keep batches separate +//! until consolidation is needed. +//! +//! ## Relationship to XArray +//! `XArray` wraps NdArray or NdArrayV and adds named dimensions and coordinate +//! labels. SuperNdArray remains a separate container. + +use std::fmt; + +use crate::enums::error::MinarrowError; +use crate::enums::shape_dim::ShapeDim; +use crate::structs::chunked::super_array::RechunkStrategy; +use crate::structs::ndarray::NdArray; +use crate::traits::concatenate::Concatenate; +use crate::traits::consolidate::Consolidate; +use crate::traits::shape::Shape; +use crate::traits::type_unions::Float; +use crate::structs::ndarray::NdArrayIter; +#[cfg(feature = "views")] +use crate::Vec64; +#[cfg(all(feature = "views", feature = "select"))] +use crate::structs::ndarray::gather_obs_impl; +#[cfg(all(feature = "views", feature = "select"))] +use crate::traits::selection::{AxisSelection, DataSelector, RowSelection}; +#[cfg(feature = "views")] +use crate::structs::views::chunked::super_ndarray_view::SuperNdArrayV; +#[cfg(feature = "views")] +use crate::structs::views::ndarray_view::NdArrayV; + +#[cfg(feature = "parallel_proc")] +use rayon::prelude::*; + +/// N-dimensional array represented by a sequence of batches. +/// +/// All batches share the same rank and non-leading dimensions. +/// Only the leading axis (axis 0) may differ between batches. +/// Global indices address the combined axis-0 range and are resolved to the +/// corresponding batch internally. +#[derive(Clone)] +pub struct SuperNdArray { + pub batches: Vec>, + ndim: usize, + inner_shape: Vec, + pub name: String, +} + +impl SuperNdArray { + /// Create an empty SuperNdArray. + pub fn new(name: impl Into) -> Self { + SuperNdArray { + batches: Vec::new(), + ndim: 0, + inner_shape: Vec::new(), + name: name.into(), + } + } + + /// Create from existing batches. Panics if shapes are incompatible. + pub fn from_batches(batches: Vec>, name: impl Into) -> Self { + Self::try_from_batches(batches, name).unwrap_or_else(|error| panic!("{error}")) + } + + /// Create from existing batches, returning an error if their shapes are + /// incompatible. + pub fn try_from_batches( + batches: Vec>, + name: impl Into, + ) -> Result { + let name = name.into(); + if batches.is_empty() { + return Ok(Self::new(name)); + } + let ndim = batches[0].ndim(); + if ndim == 0 { + return Err(MinarrowError::ShapeError { + message: "SuperNdArray batches require an axis 0".to_string(), + }); + } + let inner_shape: Vec = batches[0].shape()[1..].to_vec(); + for (index, batch) in batches.iter().enumerate().skip(1) { + if batch.ndim() != ndim { + return Err(MinarrowError::ShapeError { + message: format!( + "SuperNdArray: batch {} has rank {} but expected {}", + index, + batch.ndim(), + ndim + ), + }); + } + if &batch.shape()[1..] != inner_shape.as_slice() { + return Err(MinarrowError::ShapeError { + message: format!("SuperNdArray: batch {} inner shape mismatch", index), + }); + } + } + Ok(SuperNdArray { batches, ndim, inner_shape, name }) + } + + /// Append a batch. Validates shape compatibility. + pub fn push(&mut self, chunk: NdArray) { + assert!(chunk.ndim() > 0, "SuperNdArray batches require an axis 0"); + if self.batches.is_empty() { + self.ndim = chunk.ndim(); + self.inner_shape = chunk.shape()[1..].to_vec(); + } else { + assert_eq!( + chunk.ndim(), self.ndim, + "SuperNdArray::push: rank {} does not match expected {}", chunk.ndim(), self.ndim + ); + assert_eq!( + &chunk.shape()[1..], self.inner_shape.as_slice(), + "SuperNdArray::push: inner shape {:?} does not match expected {:?}", + &chunk.shape()[1..], self.inner_shape + ); + } + self.batches.push(chunk); + } + + // **************************************************************** + // Introspection + // **************************************************************** + + /// Number of batches. + #[inline] + pub fn n_batches(&self) -> usize { self.batches.len() } + + /// Shared rank. + #[inline] + pub fn ndim(&self) -> usize { self.ndim } + + /// Total logical elements across all batches i.e. the product of shape. + #[inline] + pub fn len(&self) -> usize { + self.batches.iter().map(|c| c.len()).sum() + } + + #[inline] + pub fn is_empty(&self) -> bool { self.len() == 0 } + + /// Total leading-axis (axis 0) observations across all batches. + #[inline] + pub fn n_obs(&self) -> usize { + self.batches.iter().map(|c| c.shape()[0]).sum() + } + + /// Dimensions shared across all batches i.e. shape[1..]. + #[inline] + pub fn inner_shape(&self) -> &[usize] { &self.inner_shape } + + /// Get batch by index. + #[inline] + pub fn batch(&self, idx: usize) -> Option<&NdArray> { self.batches.get(idx) } + + /// The constituent batches. + #[inline] + pub fn batches(&self) -> &[NdArray] { &self.batches } + + /// Consume into the constituent batches. + #[inline] + pub fn into_batches(self) -> Vec> { self.batches } + + /// Logical shape as if consolidated. Returns a temporary vec, not a + /// slice reference, since the full shape doesn't exist as a contiguous + /// field - axis 0 is the sum across batches. + pub fn shape(&self) -> Vec { + let mut s = vec![self.n_obs()]; + s.extend_from_slice(&self.inner_shape); + s + } + + /// Strides of the first batch. Strides above axis 0 scale with each + /// batch's own leading-axis length, so these describe the first batch + /// only and are not shared across batches. Use them for per-batch + /// access rather than for the consolidated array. + pub fn strides(&self) -> &[usize] { + if self.batches.is_empty() { return &[]; } + self.batches[0].strides() + } + + /// Column slice from a 2D SuperNdArray. Consolidates the column + /// across all batches into a contiguous allocation. + pub fn col(&self, c: usize) -> Vec { + assert_eq!(self.ndim, 2, "col() requires a 2D array"); + let mut result = Vec::with_capacity(self.n_obs()); + for batch in &self.batches { + result.extend_from_slice(batch.col(c)); + } + result + } + + /// Returns a zero-copy [`SuperNdArrayV`] over `[offset .. offset + len)` + /// axis-0 observations, spanning batch boundaries. Each overlapped batch + /// contributes a windowed slice view. + #[cfg(feature = "views")] + pub fn slice(&self, mut offset: usize, mut len: usize) -> SuperNdArrayV { + assert!( + offset + len <= self.n_obs(), + "SuperNdArray::slice: window [{}, {}) out of bounds (n_obs {})", + offset, offset + len, self.n_obs() + ); + + let mut slices = Vec::new(); + for batch in &self.batches { + let base_obs = batch.shape()[0]; + if offset >= base_obs { + offset -= base_obs; + continue; + } + + let take = (base_obs - offset).min(len); + let mut window_shape = vec![take]; + window_shape.extend_from_slice(&batch.shape()[1..]); + slices.push(NdArrayV::new( + batch.clone(), + offset * batch.strides()[0], + &window_shape, + batch.strides(), + )); + + len -= take; + if len == 0 { + break; + } + offset = 0; + } + + SuperNdArrayV::from_slices(slices, self.ndim, self.inner_shape.clone(), self.name.clone()) + } + + /// Zero-copy view of a single observation (axis-0 element) + /// across batch boundaries. + /// + /// Returns an (N-1)-dimensional `NdArrayV` view. For a 2D SuperNdArray + /// with shape `[n, m]`, returns a 1D view of shape `[m]`. For 3D + /// `[n, m, k]`, returns 2D `[m, k]`. Requires rank 2 or higher - for + /// scalar access on a 1D array use `get(&[i])`. + #[cfg(feature = "views")] + pub fn obs(&self, idx: usize) -> NdArrayV { + let (chunk_idx, local) = self.resolve(idx); + self.batches[chunk_idx].obs(local) + } + + /// BLAS row count (2D). + #[inline] + pub fn m(&self) -> i32 { + assert_eq!(self.ndim, 2, "m() requires a 2D array"); + self.n_obs() as i32 + } + + /// BLAS column count (2D). + #[inline] + pub fn n(&self) -> i32 { + assert_eq!(self.ndim, 2, "n() requires a 2D array"); + self.inner_shape[0] as i32 + } + + /// BLAS leading dimension of the first batch (2D). The leading dimension + /// equals each batch's own row count, so this value applies to the first + /// batch only, not to the SuperNdArray as a whole. + #[inline] + pub fn lda(&self) -> i32 { + assert_eq!(self.ndim, 2, "lda() requires a 2D array"); + if self.batches.is_empty() { return 0; } + self.batches[0].lda() + } + + // **************************************************************** + // Global element access + // **************************************************************** + + /// Get an element by global N-dimensional index. The first index is a + /// position in the combined axis-0 range and resolves to one batch. + pub fn get(&self, indices: &[usize]) -> T { + let (chunk_idx, local_row) = self.resolve(indices[0]); + let mut local = indices.to_vec(); + local[0] = local_row; + self.batches[chunk_idx].get(&local) + } + + /// Set an element by global index. Copy-on-write applies if the target + /// batch's buffer is shared. + pub fn set(&mut self, indices: &[usize], value: T) { + let (chunk_idx, local_row) = self.resolve(indices[0]); + let mut local = indices.to_vec(); + local[0] = local_row; + self.batches[chunk_idx].set(&local, value); + } + + /// Resolve a global axis-0 index to (batch_index, local_index). + /// Returns an error if the index is out of bounds. + pub(crate) fn try_resolve(&self, global: usize) -> Result<(usize, usize), MinarrowError> { + let mut remaining = global; + for (i, chunk) in self.batches.iter().enumerate() { + let n = chunk.shape()[0]; + if remaining < n { + return Ok((i, remaining)); + } + remaining -= n; + } + Err(MinarrowError::IndexError( + format!("global index {} out of bounds (n_obs {})", global, self.n_obs()) + )) + } + + /// Resolve a global axis-0 index to (batch_index, local_index). + /// Panics if the index is out of bounds. + fn resolve(&self, global: usize) -> (usize, usize) { + self.try_resolve(global).unwrap_or_else(|e| panic!("{}", e)) + } + + // **************************************************************** + // Iteration + // **************************************************************** + + /// Iterate over batches. + #[inline] + pub fn iter_batches(&self) -> std::slice::Iter<'_, NdArray> { + self.batches.iter() + } + + /// Mutable iteration over batches. + #[inline] + pub fn iter_batches_mut(&mut self) -> std::slice::IterMut<'_, NdArray> { + self.batches.iter_mut() + } + + /// Iterate values in the column-major order of the logical consolidated + /// array without materialising it. Unlike `IntoIterator`, this interleaves + /// corresponding axis-0 runs across batches. + pub fn iter_logical(&self) -> impl Iterator + '_ { + let n_runs: usize = self.inner_shape.iter().product(); + (0..n_runs).flat_map(move |run| { + self.batches.iter().flat_map(move |batch| batch.iter_axis0_run(run)) + }) + } + + /// Parallel iteration over batches. + #[cfg(feature = "parallel_proc")] + #[inline] + pub fn par_iter_batches(&self) -> rayon::slice::Iter<'_, NdArray> + where + T: Send + Sync, + { + self.batches.par_iter() + } + + /// Parallel iterator over axis-0 observations across all batches. + /// Each item is the global observation index and a zero-copy + /// `NdArrayV` view. Global indices are resolved across batch boundaries. + #[cfg(all(feature = "parallel_proc", feature = "views"))] + pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator)> + '_ + where + T: Send + Sync, + { + use rayon::prelude::*; + let n_obs = self.n_obs(); + (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i))) + } + + // **************************************************************** + // Apply + // **************************************************************** + + /// Apply a function to every logical element, returning a new + /// SuperNdArray with the same batch boundaries and name. + pub fn apply(&self, f: impl Fn(T) -> T) -> SuperNdArray { + let batches: Vec> = self.batches.iter().map(|b| b.apply(&f)).collect(); + SuperNdArray::from_batches(batches, self.name.clone()) + } + + /// Apply a function to every logical element in place, with no + /// allocation. Copy-on-write triggers per batch when views share a + /// batch's buffer. + pub fn apply_mut(&mut self, f: impl Fn(T) -> T) { + for batch in &mut self.batches { + batch.apply_mut(&f); + } + } + + /// Apply a function to every 1D lane along the given axis, collapsing + /// that axis. Each lane arrives as an [`NdArrayV`] and the closure + /// returns one value for it, mirroring [`NdArray::apply_axis`]. + /// + /// Lanes along axis 1 and above live inside single batches, so the + /// result keeps this array's batch boundaries. Lanes along axis 0 + /// cross batch boundaries, so each gathers into a contiguous buffer + /// and the result holds one batch with the trailing shape. + #[cfg(feature = "views")] + pub fn apply_axis(&self, axis: usize, mut f: impl FnMut(NdArrayV) -> T) -> SuperNdArray { + assert!(self.ndim >= 2, "apply_axis requires a 2D or higher array"); + assert!( + axis < self.ndim, + "apply_axis: axis {} out of bounds for {}D array", axis, self.ndim + ); + + if axis > 0 { + let batches: Vec> = self + .batches + .iter() + .map(|b| b.apply_axis(axis, &mut f)) + .collect(); + return SuperNdArray::from_batches(batches, self.name.clone()); + } + + // Axis-0 lanes span batches. Walk the inner positions in + // column-major order, gathering each lane batch by batch. + let n_obs = self.n_obs(); + let inner_positions: usize = self.inner_shape.iter().product(); + let mut results = Vec64::with_capacity(inner_positions); + let mut inner_idx = vec![0usize; self.inner_shape.len()]; + let mut lane_idx = vec![0usize; self.ndim]; + for _ in 0..inner_positions { + let mut lane = Vec64::with_capacity(n_obs); + lane_idx[1..].copy_from_slice(&inner_idx); + for batch in &self.batches { + let obs = batch.shape()[0]; + for i in 0..obs { + lane_idx[0] = i; + lane.push(batch.get(&lane_idx)); + } + } + let lane_arr = NdArray::from_slice(&lane, &[n_obs]); + results.push(f(lane_arr.as_view())); + + let mut carry = true; + for d in 0..inner_idx.len() { + if carry { + inner_idx[d] += 1; + if inner_idx[d] < self.inner_shape[d] { + carry = false; + } else { + inner_idx[d] = 0; + } + } + } + } + SuperNdArray::from_batches( + vec![NdArray::from_slice(&results, &self.inner_shape)], + self.name.clone(), + ) + } + + /// Apply a transformation to each batch, producing a new SuperNdArray. + /// + /// The closure receives each batch and returns a transformed batch, + /// mirroring `Table::apply_cols` with the batch as the unit of work. + /// Returned batches must share rank and trailing shape. + pub fn apply_batches( + &self, + mut f: impl FnMut(&NdArray) -> Result, E>, + ) -> Result, E> { + let batches = self + .batches + .iter() + .map(|b| f(b)) + .collect::, E>>()?; + Ok(SuperNdArray::from_batches(batches, self.name.clone())) + } + + /// Apply an in-place transformation to each batch, with no allocation + /// beyond what the closure itself performs. + pub fn apply_batches_mut( + &mut self, + mut f: impl FnMut(&mut NdArray) -> Result<(), E>, + ) -> Result<(), E> { + for batch in &mut self.batches { + f(batch)?; + } + Ok(()) + } + + // **************************************************************** + // Changing batch boundaries + // **************************************************************** + + /// Redistribute data across batches. + /// + /// - `Count(n)` - uniform batches of n leading-axis elements + /// - `Auto` - default batch length of 8192 leading-axis elements + /// - `Memory(bytes)` - target byte size per batch (requires `size` feature) + pub fn rechunk(&mut self, strategy: RechunkStrategy) -> Result<(), MinarrowError> { + if self.batches.is_empty() || self.n_obs() == 0 { + return Ok(()); + } + + let chunk_size = match strategy { + RechunkStrategy::Count(size) => { + if size == 0 { + return Err(MinarrowError::IndexError( + "rechunk: chunk size must be > 0".to_string(), + )); + } + size + } + RechunkStrategy::Auto => 8192, + #[cfg(feature = "size")] + RechunkStrategy::Memory(target_bytes) => { + use crate::traits::byte_size::ByteSize; + if target_bytes == 0 { + return Err(MinarrowError::IndexError( + "rechunk: target bytes must be > 0".to_string(), + )); + } + // Estimate observations per batch from the first batch's byte density. + let first = &self.batches[0]; + let bytes_per_row = first.est_bytes().max(1) / first.shape()[0].max(1); + (target_bytes / bytes_per_row.max(1)).max(1) + } + }; + + // Capture state before moving self + let total_obs = self.n_obs(); + let ndim = self.ndim; + let inner_shape = self.inner_shape.clone(); + let name = self.name.clone(); + + // Consolidate then re-split + let old = std::mem::replace(self, SuperNdArray::new(&name)); + let consolidated = old.consolidate(); + let strides = consolidated.strides().to_vec(); + let buf = consolidated.as_slice(); + + let mut new_batches = Vec::with_capacity(total_obs.div_ceil(chunk_size)); + let mut row_offset = 0; + + while row_offset < total_obs { + let n = (total_obs - row_offset).min(chunk_size); + let mut chunk_shape = vec![n]; + chunk_shape.extend_from_slice(&inner_shape); + + if ndim <= 1 { + let chunk = NdArray::from_slice(&buf[row_offset..row_offset + n], &chunk_shape); + new_batches.push(chunk); + } else { + let mut chunk = NdArray::new(&chunk_shape); + let dst_stride = chunk.strides()[1]; + let n_cols: usize = inner_shape.iter().product::().max(1); + let dst = chunk.as_mut_slice(); + for c in 0..n_cols { + let src_start = c * strides[1] + row_offset; + let dst_start = c * dst_stride; + dst[dst_start..dst_start + n] + .copy_from_slice(&buf[src_start..src_start + n]); + } + new_batches.push(chunk); + } + row_offset += n; + } + + *self = SuperNdArray { ndim, inner_shape, batches: new_batches, name }; + Ok(()) + } +} + +// *** Axis selection: snd.s(nd![1..4, 2]) ************************* + +/// Selection across every axis at once. The axis-0 selection windows +/// across batch boundaries, and trailing-axis selections narrow each +/// slice. Zero-copy. Delegates to the full-window [`SuperNdArrayV`]. +/// +/// An axis-0 single index keeps the dimension as a one-observation +/// window, since a SuperNdArrayV cannot collapse its shared axis - use +/// `obs` to collapse a single observation. Trailing single indices +/// collapse their dimensions as usual. +#[cfg(all(feature = "views", feature = "select"))] +impl AxisSelection for SuperNdArray { + type View = SuperNdArrayV; + + fn s(&self, selection: &[&dyn DataSelector]) -> SuperNdArrayV { + self.slice(0, self.n_obs()).s(selection) + } + + fn get_axis_count(&self) -> usize { + self.ndim + } +} + +// *** Row selection: snd.r(0..10) ********************************* + +/// Axis-0 observation selection across batch boundaries. Contiguous +/// ranges return a zero-copy [`SuperNdArrayV`] window. Index arrays +/// gather the selected observations into one owned batch wrapped in a +/// single-slice view. +#[cfg(all(feature = "views", feature = "select"))] +impl RowSelection for SuperNdArray { + type View = SuperNdArrayV; + + fn r(&self, selection: S) -> SuperNdArrayV { + if self.batches.is_empty() { + return SuperNdArrayV::from_slices( + Vec::new(), + self.ndim, + self.inner_shape.clone(), + self.name.clone(), + ); + } + let indices = selection.resolve_indices(self.n_obs()); + if selection.is_contiguous() { + let start = indices.first().copied().unwrap_or(0); + return self.slice(start, indices.len()); + } + let gathered = gather_obs_impl( + &indices, + &self.shape(), + Some(self.name.clone()), + |idx| self.get(idx), + ); + SuperNdArrayV::from_slices( + vec![NdArrayV::from_ndarray(gathered)], + self.ndim, + self.inner_shape.clone(), + self.name.clone(), + ) + } + + fn get_row_count(&self) -> usize { + self.n_obs() + } +} + +// **************************************************************** +// IntoIterator - element-wise iteration across batches +// **************************************************************** + +/// Iterating a SuperNdArray walks each batch in sequence, with column-major +/// order inside each batch. Use [`SuperNdArray::iter_logical`] when values +/// must follow the column-major order of the consolidated logical array. +impl<'a, T: Float> IntoIterator for &'a SuperNdArray { + type Item = T; + type IntoIter = SuperNdArrayIter<'a, T>; + + fn into_iter(self) -> SuperNdArrayIter<'a, T> { + SuperNdArrayIter { + batches: &self.batches, + chunk_idx: 0, + inner: None, + } + } +} + +/// Iterator over the stored batches in sequence. +pub struct SuperNdArrayIter<'a, T> { + batches: &'a [NdArray], + chunk_idx: usize, + inner: Option>, +} + +impl<'a, T: Float> Iterator for SuperNdArrayIter<'a, T> { + type Item = T; + + fn next(&mut self) -> Option { + loop { + if let Some(ref mut it) = self.inner { + if let Some(v) = it.next() { + return Some(v); + } + } + // Continue with the next batch. + if self.chunk_idx >= self.batches.len() { + return None; + } + self.inner = Some((&self.batches[self.chunk_idx]).into_iter()); + self.chunk_idx += 1; + } + } + + fn size_hint(&self) -> (usize, Option) { + let inner_remaining = self.inner.as_ref().map_or(0, |it| it.len()); + let remaining_batches: usize = self.batches[self.chunk_idx..] + .iter() + .map(|c| c.len()) + .sum(); + let total = inner_remaining + remaining_batches; + (total, Some(total)) + } +} + +impl<'a, T: Float> ExactSizeIterator for SuperNdArrayIter<'a, T> {} + +// **************************************************************** +// Consolidate +// **************************************************************** + +impl Consolidate for SuperNdArray { + type Output = NdArray; + + /// Materialise into a single contiguous column-major [`NdArray`]. + /// Each source batch reads at its own strides and writes straight + /// into the result, with unit-stride axis-0 runs taking the + /// memcpy path. + fn consolidate(self) -> NdArray { + if self.batches.is_empty() { + let mut result = NdArray::new(&[0]); + result.name = Some(self.name); + return result; + } + if self.batches.len() == 1 { + let mut result = self.batches.into_iter().next().unwrap(); + if !result.is_contiguous() { + result = result.to_contiguous(); + } + result.name = Some(self.name); + return result; + } + + let total_axis0: usize = self.batches.iter().map(|c| c.shape()[0]).sum(); + let mut full_shape = vec![total_axis0]; + full_shape.extend_from_slice(&self.inner_shape); + + let mut result = NdArray::new(&full_shape); + + if self.ndim == 1 { + let dst = result.as_mut_slice(); + let mut pos = 0; + for chunk in &self.batches { + let n = chunk.shape()[0]; + let s0 = chunk.strides()[0]; + let src = chunk.as_slice(); + if s0 == 1 { + dst[pos..pos + n].copy_from_slice(&src[..n]); + } else { + for i in 0..n { + dst[pos + i] = src[i * s0]; + } + } + pos += n; + } + } else { + // The result is contiguous column-major, so column c starts at + // c * strides[1] for c in 0..product(inner_shape), and the + // axis-0 rows from every batch interleave one column at a time. + // Column c's base offset within each source batch decomposes + // across that batch's own outer dimensions; for a contiguous + // batch this reduces to c * strides[1]. + let dst_stride = result.strides()[1]; + let n_cols: usize = self.inner_shape.iter().product(); + let dst = result.as_mut_slice(); + let mut row_offset = 0; + for chunk in &self.batches { + let chunk_obs = chunk.shape()[0]; + let s0 = chunk.strides()[0]; + let src = chunk.as_slice(); + for c in 0..n_cols { + let mut src_start = 0; + let mut rem = c; + for d in 1..chunk.ndim() { + src_start += (rem % chunk.shape()[d]) * chunk.strides()[d]; + rem /= chunk.shape()[d]; + } + let dst_start = c * dst_stride + row_offset; + if s0 == 1 { + dst[dst_start..dst_start + chunk_obs] + .copy_from_slice(&src[src_start..src_start + chunk_obs]); + } else { + for i in 0..chunk_obs { + dst[dst_start + i] = src[src_start + i * s0]; + } + } + } + row_offset += chunk_obs; + } + } + + result.name = Some(self.name); + result + } +} + +// **************************************************************** +// Trait implementations +// **************************************************************** + +/// Logical equality. Two SuperNdArrays are equal when they share the same +/// rank, trailing shape, and observation count, and hold the same values in +/// logical order. Batch boundaries and the name do not affect equality. +/// +/// The comparison walks one logical column at a time, chaining each side's +/// batches independently, so different batch divisions still line +/// up without materialising either side. Each batch reads at its own +/// strides, so any layout compares in place. +impl PartialEq for SuperNdArray { + fn eq(&self, other: &Self) -> bool { + if self.ndim != other.ndim + || self.inner_shape != other.inner_shape + || self.n_obs() != other.n_obs() + { + return false; + } + let n_cols: usize = self.inner_shape.iter().product(); + for c in 0..n_cols { + let lhs = self.batches.iter().flat_map(|b| { + let obs = b.shape()[0]; + let s0 = b.strides()[0]; + // Column c's base offset decomposes across the outer dims, + // which for a contiguous batch reduces to c * strides[1]. + let mut off = 0; + let mut rem = c; + for d in 1..b.ndim() { + off += (rem % b.shape()[d]) * b.strides()[d]; + rem /= b.shape()[d]; + } + let buf = b.as_slice(); + (0..obs).map(move |i| buf[off + i * s0]) + }); + let rhs = other.batches.iter().flat_map(|b| { + let obs = b.shape()[0]; + let s0 = b.strides()[0]; + let mut off = 0; + let mut rem = c; + for d in 1..b.ndim() { + off += (rem % b.shape()[d]) * b.strides()[d]; + rem /= b.shape()[d]; + } + let buf = b.as_slice(); + (0..obs).map(move |i| buf[off + i * s0]) + }); + if !lhs.eq(rhs) { + return false; + } + } + true + } +} + +impl Shape for SuperNdArray { + fn shape(&self) -> ShapeDim { + let obs = self.n_obs(); + match self.ndim { + 0 | 1 => ShapeDim::Rank1(obs), + 2 => ShapeDim::Rank2 { rows: obs, cols: self.inner_shape[0] }, + _ => { + let mut full = vec![obs]; + full.extend_from_slice(&self.inner_shape); + ShapeDim::RankN(full) + } + } + } +} + +impl Concatenate for SuperNdArray { + fn concat(mut self, other: Self) -> Result { + if self.batches.is_empty() { return Ok(other); } + if other.batches.is_empty() { return Ok(self); } + if self.ndim != other.ndim { + return Err(MinarrowError::IncompatibleTypeError { + from: "SuperNdArray", to: "SuperNdArray", + message: Some(format!("rank {} vs {}", self.ndim, other.ndim)), + }); + } + if self.inner_shape != other.inner_shape { + return Err(MinarrowError::IncompatibleTypeError { + from: "SuperNdArray", to: "SuperNdArray", + message: Some(format!( + "inner shape {:?} vs {:?}", self.inner_shape, other.inner_shape + )), + }); + } + self.batches.extend(other.batches); + Ok(self) + } +} + +impl fmt::Debug for SuperNdArray { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, "SuperNdArray '{}': {} batches, {}D, shape {:?}, {} elements", + self.name, self.n_batches(), self.ndim, self.shape(), self.len() + ) + } +} + +// **************************************************************** +// Tests +// **************************************************************** + +#[cfg(test)] +mod tests { + use super::*; + use crate::Buffer; + #[cfg(all(feature = "views", feature = "select"))] + use crate::nd; + + // *** Row selection and apply ************************************* + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn axis_selection_windows_and_narrows() { + // col0 = [1..5], col1 = [10..50] across two batches. + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2]), + NdArray::from_slice(&[3.0, 4.0, 5.0, 30.0, 40.0, 50.0], &[3, 2]), + ], + "data", + ); + // Axis-0 range crosses the batch boundary, trailing index collapses. + let v = snd.s(nd![1..4, 1]); + assert_eq!(v.ndim(), 1); + assert_eq!(v.n_obs(), 3); + assert_eq!(v.get(&[0]), 20.0); + assert_eq!(v.get(&[2]), 40.0); + // Axis-0 single index keeps the dimension as a one-observation window. + let one = snd.s(nd![3, 0..2]); + assert_eq!(one.n_obs(), 1); + assert_eq!(one.get(&[0, 0]), 4.0); + assert_eq!(one.get(&[0, 1]), 40.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn row_selection_contiguous_crosses_batches() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2]), + NdArray::from_slice(&[3.0, 4.0, 5.0, 30.0, 40.0, 50.0], &[3, 2]), + ], + "data", + ); + let v = snd.r(1..4); + assert_eq!(v.n_slices(), 2); + assert_eq!(v.n_obs(), 3); + assert_eq!(v.get(&[0, 0]), 2.0); + assert_eq!(v.get(&[2, 1]), 40.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn row_selection_gathers_across_batches() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0, 4.0], &[2]), + ], + "data", + ); + let v = snd.r(&[3, 0]); + assert_eq!(v.n_slices(), 1); + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals, vec![4.0, 1.0]); + } + + #[test] + fn apply_preserves_chunking() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "stream", + ); + let out = snd.apply(|x| x * 10.0); + assert_eq!(out.n_batches(), 2); + assert_eq!(out.name, "stream"); + let vals: Vec = (&out).into_iter().collect(); + assert_eq!(vals, vec![10.0, 20.0, 30.0]); + } + + #[test] + fn apply_mut_in_place() { + let mut snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "stream", + ); + snd.apply_mut(|x| x + 1.0); + assert_eq!(snd.n_batches(), 2); + let vals: Vec = (&snd).into_iter().collect(); + assert_eq!(vals, vec![2.0, 3.0, 4.0]); + } + + #[cfg(feature = "views")] + #[test] + fn apply_axis_zero_crosses_batches() { + // col0 = [1, 2, 3, 4], col1 = [10, 20, 30, 40] across two batches. + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2]), + NdArray::from_slice(&[3.0, 4.0, 30.0, 40.0], &[2, 2]), + ], + "data", + ); + let sums = snd.apply_axis(0, |lane| (&lane).into_iter().sum()); + assert_eq!(sums.n_batches(), 1); + assert_eq!(sums.shape(), vec![2]); + assert_eq!(sums.get(&[0]), 10.0); + assert_eq!(sums.get(&[1]), 100.0); + } + + #[cfg(feature = "views")] + #[test] + fn apply_axis_inner_preserves_chunking() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2]), + NdArray::from_slice(&[3.0, 30.0], &[1, 2]), + ], + "data", + ); + // Sum each row lane (axis 1) - per-batch, boundaries kept. + let row_sums = snd.apply_axis(1, |lane| (&lane).into_iter().sum()); + assert_eq!(row_sums.n_batches(), 2); + assert_eq!(row_sums.shape(), vec![3]); + assert_eq!(row_sums.get(&[0]), 11.0); + assert_eq!(row_sums.get(&[1]), 22.0); + assert_eq!(row_sums.get(&[2]), 33.0); + } + + #[test] + fn apply_batches_transforms_each_batch() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "stream", + ); + let out = snd + .apply_batches(|b| Ok::<_, MinarrowError>(b.apply(|x| x - 1.0))) + .unwrap(); + assert_eq!(out.n_batches(), 2); + let vals: Vec = (&out).into_iter().collect(); + assert_eq!(vals, vec![0.0, 1.0, 2.0]); + } + + #[test] + fn apply_batches_propagates_errors() { + let snd = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0], &[1])], + "stream", + ); + let result = snd.apply_batches(|_| { + Err::, MinarrowError>(MinarrowError::KernelError(Some( + "boom".to_string(), + ))) + }); + assert!(result.is_err()); + } + + #[test] + fn apply_batches_mut_in_place() { + let mut snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "stream", + ); + snd.apply_batches_mut(|b| { + b.apply_mut(|x| x * 3.0); + Ok::<_, MinarrowError>(()) + }) + .unwrap(); + let vals: Vec = (&snd).into_iter().collect(); + assert_eq!(vals, vec![3.0, 6.0, 9.0]); + } + + #[test] + fn empty() { + let snd = SuperNdArray::::new("empty"); + assert!(snd.is_empty()); + assert_eq!(snd.n_batches(), 0); + } + + #[test] + fn from_batches_1d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let b = NdArray::from_slice(&[4.0, 5.0], &[2]); + let snd = SuperNdArray::from_batches(vec![a, b], "test"); + assert_eq!(snd.n_batches(), 2); + assert_eq!(snd.len(), 5); + assert_eq!(snd.n_obs(), 5); + assert_eq!(snd.ndim(), 1); + } + + #[test] + fn from_batches_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]); + let snd = SuperNdArray::from_batches(vec![a, b], "data"); + assert_eq!(snd.n_batches(), 2); + assert_eq!(snd.len(), 10); // 2*2 + 3*2 = 10 total elements + assert_eq!(snd.n_obs(), 5); // 2 + 3 = 5 leading-axis observations + assert_eq!(snd.inner_shape(), &[2]); + } + + #[test] + fn push_and_validate() { + let mut snd = SuperNdArray::new("stream"); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])); + snd.push(NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2])); + assert_eq!(snd.n_batches(), 2); + assert_eq!(snd.len(), 8); // 2*2 + 2*2 = 8 total elements + assert_eq!(snd.n_obs(), 4); + } + + #[test] + #[should_panic(expected = "rank")] + fn push_rank_mismatch() { + let mut snd = SuperNdArray::new("bad"); + snd.push(NdArray::from_slice(&[1.0, 2.0], &[2])); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])); + } + + #[test] + #[should_panic(expected = "inner shape")] + fn push_trailing_shape_mismatch() { + let mut snd = SuperNdArray::new("bad"); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])); + } + + // *** Global element access *************************************** + + #[test] + fn global_get_1d() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]), + NdArray::from_slice(&[40.0, 50.0], &[2]), + ], "test"); + assert_eq!(snd.get(&[0]), 10.0); + assert_eq!(snd.get(&[2]), 30.0); + assert_eq!(snd.get(&[3]), 40.0); + assert_eq!(snd.get(&[4]), 50.0); + } + + #[test] + fn global_get_2d() { + // Batch 0: 2x2, col0=[1,2], col1=[3,4]. + // Batch 1: 3x2, col0=[5,6,7], col1=[8,9,10]. + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "data"); + // Global row 0, col 0 -> batch 0 row 0, col 0. + assert_eq!(snd.get(&[0, 0]), 1.0); + // Global row 1, col 1 -> batch 0 row 1, col 1. + assert_eq!(snd.get(&[1, 1]), 4.0); + // Global row 2, col 0 -> batch 1 row 0, col 0. + assert_eq!(snd.get(&[2, 0]), 5.0); + // Global row 4, col 1 -> batch 1 row 2, col 1. + assert_eq!(snd.get(&[4, 1]), 10.0); + } + + #[test] + fn global_set_2d() { + let mut snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + ], "mut"); + snd.set(&[3, 1], 99.0); + assert_eq!(snd.get(&[3, 1]), 99.0); + } + + // *** Iteration *************************************************** + + #[test] + fn iter_1d_across_batches() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]), + NdArray::from_slice(&[4.0, 5.0], &[2]), + ], "test"); + let vals: Vec = (&snd).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + } + + #[test] + fn iter_2d_across_batches() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + ], "data"); + let vals: Vec = (&snd).into_iter().collect(); + // Batch 0 is [1,2,3,4] and batch 1 is [5,6,7,8] in column-major order. + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); + + let logical: Vec = snd.iter_logical().collect(); + assert_eq!(logical, vec![1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 7.0, 8.0]); + assert_eq!( + logical, + snd.clone().consolidate().into_iter().collect::>() + ); + } + + #[test] + fn iter_exact_size() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]), + NdArray::from_slice(&[4.0, 5.0], &[2]), + ], "test"); + let iter = (&snd).into_iter(); + assert_eq!(iter.len(), 5); + } + + // *** Consolidate ************************************************* + + #[test] + fn consolidate_1d() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]), + NdArray::from_slice(&[4.0, 5.0], &[2]), + ], "test"); + let result = snd.consolidate(); + assert_eq!(result.shape(), &[5]); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + } + + #[test] + fn consolidate_2d() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "data"); + let result = snd.consolidate(); + assert_eq!(result.shape(), &[5, 2]); + assert_eq!(result.col(0), &[1.0, 2.0, 5.0, 6.0, 7.0]); + assert_eq!(result.col(1), &[3.0, 4.0, 8.0, 9.0, 10.0]); + } + + #[test] + fn consolidate_single_chunk() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let snd = SuperNdArray::from_batches(vec![a.clone()], "one"); + let result = snd.consolidate(); + // The single-batch shortcut carries the SuperNdArray's name. + assert_eq!(result.name.as_deref(), Some("one")); + let mut expected = a; + expected.name = Some("one".to_string()); + assert_eq!(result, expected); + } + + #[test] + fn batch_accessors() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "ticks", + ); + assert_eq!(snd.batches().len(), 2); + assert_eq!(snd.batch(1).unwrap().get(&[0]), 3.0); + assert!(snd.batch(2).is_none()); + let batches = snd.into_batches(); + assert_eq!(batches.len(), 2); + } + + #[cfg(feature = "views")] + #[test] + fn view_carries_name() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0, 4.0], &[2]), + ], + "ticks", + ); + let window = snd.slice(1, 2); + assert_eq!(window.name(), "ticks"); + assert_eq!(window.consolidate().name.as_deref(), Some("ticks")); + } + + #[test] + fn consolidate_3d() { + // Batch A is [2,2,2] with logical column-major values 1..=8; + // batch B is [1,2,2] with logical column-major values 9..=12. + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 2, 2]); + let b = NdArray::from_slice(&[9.0, 10.0, 11.0, 12.0], &[1, 2, 2]); + let snd = SuperNdArray::from_batches(vec![a, b], "cube"); + let c = snd.consolidate(); + assert_eq!(c.shape(), &[3, 2, 2]); + // Axis-0 rows interleave one higher-dimensional slice at a time. + assert_eq!(c.get(&[0, 0, 0]), 1.0); + assert_eq!(c.get(&[1, 0, 0]), 2.0); + assert_eq!(c.get(&[2, 0, 0]), 9.0); + assert_eq!(c.get(&[0, 1, 0]), 3.0); + assert_eq!(c.get(&[2, 1, 0]), 10.0); + assert_eq!(c.get(&[1, 1, 1]), 8.0); + assert_eq!(c.get(&[2, 0, 1]), 11.0); + assert_eq!(c.get(&[2, 1, 1]), 12.0); + } + + #[test] + fn eq_ignores_chunking_and_name() { + // Same logical [3,2] values with different batch divisions and names. + let single = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2])], + "single", + ); + let split = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 4.0], &[1, 2]), + NdArray::from_slice(&[2.0, 3.0, 5.0, 6.0], &[2, 2]), + ], + "split", + ); + assert_eq!(single, split); + + // A changed value breaks equality. + let different = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 99.0], &[3, 2])], + "single", + ); + assert_ne!(single, different); + } + + #[test] + fn eq_chunk_invariant_3d() { + // Same logical [3,2,2] values: one batch versus a [1,2,2] and + // [2,2,2] division, exercising the rank-3 column walk in equality. + let whole = SuperNdArray::from_batches( + vec![NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0], + &[3, 2, 2], + )], + "whole", + ); + let split = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 4.0, 7.0, 10.0], &[1, 2, 2]), + NdArray::from_slice(&[2.0, 3.0, 5.0, 6.0, 8.0, 9.0, 11.0, 12.0], &[2, 2, 2]), + ], + "split", + ); + assert_eq!(whole, split); + } + + // *** Rechunk ***************************************************** + + #[test] + fn rechunk_count() { + let mut snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0], &[5]), + ], "test"); + snd.rechunk(RechunkStrategy::Count(2)).unwrap(); + assert_eq!(snd.n_batches(), 3); // 2 + 2 + 1 + assert_eq!(snd.len(), 5); + // Verify data integrity + let vals: Vec = (&snd).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + } + + #[test] + fn rechunk_2d() { + let mut snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]), + ], "test"); + snd.rechunk(RechunkStrategy::Count(2)).unwrap(); + assert_eq!(snd.n_batches(), 2); // 2 + 1 + assert_eq!(snd.len(), 6); // 3*2 = 6 total elements + assert_eq!(snd.n_obs(), 3); + // Verify data + assert_eq!(snd.get(&[0, 0]), 1.0); + assert_eq!(snd.get(&[2, 1]), 6.0); + } + + #[test] + fn rechunk_auto() { + let data: Vec = (0..100).map(|x| x as f64).collect(); + let mut snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&data, &[100]), + ], "big"); + snd.rechunk(RechunkStrategy::Auto).unwrap(); + // Auto uses 8192 observations, so 100 observations remain one batch. + assert_eq!(snd.n_batches(), 1); + } + + // *** Concat ****************************************************** + + #[test] + fn concat_super_ndarrays() { + let a = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0], &[2])], "a" + ); + let b = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[3.0, 4.0], &[2])], "b" + ); + let c = a.concat(b).unwrap(); + assert_eq!(c.n_batches(), 2); + assert_eq!(c.len(), 4); + } + + // *** Shape ******************************************************* + + #[test] + fn shape_trait() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + ], "test"); + assert_eq!(Shape::shape(&snd), ShapeDim::Rank2 { rows: 4, cols: 2 }); + } + + #[test] + fn shape_method() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "test"); + assert_eq!(snd.shape(), vec![5, 2]); + } + + // *** Strided batches ********************************************* + + #[test] + fn consolidate_non_contiguous_single_batch() { + // Row-major strides [2, 1] on shape [2, 2] give logical + // col0 = [1, 3] and col1 = [2, 4]. + let nd = NdArray::from_buffer( + Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]), + &[2, 2], + &[2, 1], + ); + assert!(!nd.is_contiguous()); + let snd = SuperNdArray::from_batches(vec![nd], "strided"); + let result = snd.consolidate(); + assert!(result.is_contiguous()); + assert_eq!(result.shape(), &[2, 2]); + assert_eq!(result.col(0), &[1.0, 3.0]); + assert_eq!(result.col(1), &[2.0, 4.0]); + } + + #[test] + fn consolidate_mixed_layout_batches() { + let compact = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[2, 2]); + // Row-major batch holding logical col0 = [1, 3], col1 = [2, 4]. + let strided = NdArray::from_buffer( + Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]), + &[2, 2], + &[2, 1], + ); + let snd = SuperNdArray::from_batches(vec![compact, strided], "mixed"); + let result = snd.consolidate(); + assert_eq!(result.shape(), &[4, 2]); + assert_eq!(result.col(0), &[10.0, 20.0, 1.0, 3.0]); + assert_eq!(result.col(1), &[30.0, 40.0, 2.0, 4.0]); + } + + #[test] + fn eq_strided_batch_matches_contiguous() { + // Row-major batch holding logical col0 = [1, 3], col1 = [2, 4], + // compared against the same values divided into one-row batches. + let strided = SuperNdArray::from_batches( + vec![NdArray::from_buffer( + Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]), + &[2, 2], + &[2, 1], + )], + "strided", + ); + let compact = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[1, 2]), + NdArray::from_slice(&[3.0, 4.0], &[1, 2]), + ], + "compact", + ); + assert_eq!(strided, compact); + + let different = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 3.0, 2.0, 99.0], &[2, 2])], + "different", + ); + assert_ne!(strided, different); + } + + // *** Bounds and observation access ******************************* + + #[cfg(feature = "views")] + #[test] + fn obs_resolves_across_chunk_boundary() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "data"); + // Global observation 2 is the second batch's first row. + let o = snd.obs(2); + assert_eq!(o.shape(), &[2]); + assert_eq!(o.get(&[0]), 5.0); + assert_eq!(o.get(&[1]), 8.0); + } + + #[cfg(feature = "views")] + #[test] + #[should_panic(expected = "global index 5 out of bounds (n_obs 5)")] + fn obs_out_of_bounds() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "data"); + snd.obs(5); + } + + #[test] + #[should_panic(expected = "global index 9 out of bounds (n_obs 5)")] + fn get_out_of_bounds() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]), + NdArray::from_slice(&[4.0, 5.0], &[2]), + ], "test"); + snd.get(&[9]); + } + + #[test] + #[should_panic(expected = "global index 4 out of bounds (n_obs 4)")] + fn set_out_of_bounds() { + let mut snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + ], "mut"); + snd.set(&[4, 0], 1.0); + } + + #[cfg(feature = "views")] + #[test] + #[should_panic(expected = "window [3, 7) out of bounds (n_obs 5)")] + fn slice_beyond_n_obs() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]), + NdArray::from_slice(&[4.0, 5.0], &[2]), + ], "test"); + snd.slice(3, 4); + } + + #[test] + fn col_concatenates_across_batches() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]), + ], "data"); + assert_eq!(snd.col(0), vec![1.0, 2.0, 5.0, 6.0, 7.0]); + assert_eq!(snd.col(1), vec![3.0, 4.0, 8.0, 9.0, 10.0]); + } + + // *** Construction edge cases ************************************* + + #[test] + fn from_batches_empty_then_push_adopts_shape() { + let mut snd = SuperNdArray::::from_batches(vec![], "fresh"); + assert_eq!(snd.n_batches(), 0); + assert_eq!(snd.ndim(), 0); + assert!(snd.is_empty()); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])); + assert_eq!(snd.ndim(), 2); + assert_eq!(snd.inner_shape(), &[2]); + assert_eq!(snd.n_obs(), 2); + } + + #[test] + fn try_from_batches_empty_then_push_adopts_shape() { + let mut snd = SuperNdArray::::try_from_batches(vec![], "fresh").unwrap(); + assert_eq!(snd.n_batches(), 0); + assert_eq!(snd.ndim(), 0); + snd.push(NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])); + assert_eq!(snd.inner_shape(), &[2]); + } + + #[test] + fn try_from_batches_rejects_rank_zero() { + let error = SuperNdArray::try_from_batches( + vec![NdArray::from_slice(&[1.0], &[])], + "bad", + ) + .unwrap_err(); + assert_eq!( + error, + MinarrowError::ShapeError { + message: "SuperNdArray batches require an axis 0".to_string(), + } + ); + } + + #[test] + fn try_from_batches_rejects_incompatible_shapes() { + let rank_error = SuperNdArray::try_from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + ], + "bad", + ) + .unwrap_err(); + assert_eq!( + rank_error, + MinarrowError::ShapeError { + message: "SuperNdArray: batch 1 has rank 2 but expected 1".to_string(), + } + ); + + let inner_shape_error = SuperNdArray::try_from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]), + ], + "bad", + ) + .unwrap_err(); + assert_eq!( + inner_shape_error, + MinarrowError::ShapeError { + message: "SuperNdArray: batch 1 inner shape mismatch".to_string(), + } + ); + } + + #[test] + #[should_panic(expected = "has rank 2 but expected 1")] + fn from_batches_rank_mismatch() { + SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + ], "bad"); + } + + #[test] + #[should_panic(expected = "batch 1 inner shape mismatch")] + fn from_batches_inner_shape_mismatch() { + SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]), + ], "bad"); + } + + #[test] + fn concat_rank_mismatch_errors() { + let a = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0], &[2])], "a" + ); + let b = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])], "b" + ); + let err = a.concat(b).unwrap_err(); + assert!(format!("{}", err).contains("rank 1 vs 2")); + } + + #[test] + fn concat_inner_shape_mismatch_errors() { + let a = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2])], "a" + ); + let b = SuperNdArray::from_batches( + vec![NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3])], "b" + ); + let err = a.concat(b).unwrap_err(); + assert!(format!("{}", err).contains("inner shape [2] vs [3]")); + } + + // *** Rechunk and empty batches *********************************** + + #[test] + fn rechunk_3d_count() { + let data: Vec = (1..=16).map(|x| x as f64).collect(); + let mut snd = SuperNdArray::from_batches( + vec![NdArray::from_slice(&data, &[4, 2, 2])], "cube" + ); + snd.rechunk(RechunkStrategy::Count(3)).unwrap(); + assert_eq!(snd.n_batches(), 2); // 3 + 1 + assert_eq!(snd.batch(0).unwrap().shape(), &[3, 2, 2]); + assert_eq!(snd.batch(1).unwrap().shape(), &[1, 2, 2]); + assert_eq!(snd.get(&[0, 0, 0]), 1.0); + assert_eq!(snd.get(&[2, 1, 0]), 7.0); + assert_eq!(snd.get(&[3, 0, 1]), 12.0); + assert_eq!(snd.get(&[3, 1, 1]), 16.0); + } + + #[test] + fn zero_observation_batch() { + let snd = SuperNdArray::from_batches(vec![ + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + NdArray::from_slice(&[], &[0, 2]), + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + ], "gappy"); + assert_eq!(snd.n_obs(), 4); + let vals: Vec = (&snd).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); + let result = snd.consolidate(); + assert_eq!(result.shape(), &[4, 2]); + assert_eq!(result.col(0), &[1.0, 2.0, 5.0, 6.0]); + assert_eq!(result.col(1), &[3.0, 4.0, 7.0, 8.0]); + } +} diff --git a/src/structs/matrix.rs b/src/structs/matrix.rs index 45e7176..f25435a 100644 --- a/src/structs/matrix.rs +++ b/src/structs/matrix.rs @@ -90,7 +90,7 @@ const ALIGN_ELEMS: usize = 64 / std::mem::size_of::(); // 8 /// Round up to next multiple of ALIGN_ELEMS for 64-byte column alignment. #[inline] -const fn aligned_stride(n_rows: usize) -> usize { +pub(crate) const fn aligned_stride(n_rows: usize) -> usize { (n_rows + ALIGN_ELEMS - 1) & !(ALIGN_ELEMS - 1) } diff --git a/src/structs/ndarray.rs b/src/structs/ndarray.rs new file mode 100644 index 0000000..d8cd603 --- /dev/null +++ b/src/structs/ndarray.rs @@ -0,0 +1,3498 @@ +//! # `NdArray` +//! +//! An N-dimensional container for contiguous `f32` or `f64` data. +//! +//! `NdArray` is Minarrow's general-purpose interchange container for numeric +//! data received from files, sensors, networks, Python, DLPack, or other +//! runtimes. It supports construction, indexing, views, selection, iteration, +//! conversion, and transfer to external numerical systems. +//! +//! It is not intended to provide a complete numerical-computing runtime. +//! Computational workloads should be delegated to the application's preferred +//! kernels, BLAS implementation, or machine-learning framework. +//! +//! ## Shape and element type +//! +//! `NdArray` is generic over `T: Float`, supporting `f32` and `f64`. Rank is +//! determined at runtime. Shapes with one to five dimensions are stored inline, +//! avoiding a heap allocation for the common case. +//! +//! Shape semantics follow standard tensor conventions: +//! +//! - `&[]` is a rank-zero array containing one logical value. +//! - `[0]` is a rank-one array containing no values. +//! +//! ## Memory layout +//! +//! Data is stored in compact column-major order without padding between +//! dimensions. The complete logical array therefore occupies one contiguous +//! buffer, and its shape and strides describe every stored element without +//! hidden gaps. +//! +//! The allocation start is aligned to 64 bytes by [`Vec64`], allowing efficient +//! SIMD operations over the flattened buffer. +//! +//! Per-column padding for CPU-oriented BLAS and LAPACK access is provided by +//! [`Matrix`] instead. Converting an `NdArray` with `to_matrix` re-lays the data +//! into that padded representation when required. +//! +//! NumPy and PyTorch can consume the column-major strides directly. Consumers +//! that require row-major C-contiguous storage must create a contiguous copy, +//! such as with `np.ascontiguousarray(...)` or `.contiguous()`. +//! +//! ## Null values +//! +//! `NdArray` does not store a null mask. When converting nullable floating-point +//! data, nulls are represented as `NaN`. The resulting array cannot distinguish +//! a source null from an ordinary `NaN`; interpretation is left to the consuming +//! kernel or framework. +//! +//! ## Interoperability +//! +//! - [`Matrix`] to a two-dimensional `NdArray`: zero-copy for `f64`; the padded +//! column stride is preserved. +//! - Two-dimensional `NdArray` to [`Matrix`]: re-lays data into the padded column +//! format for `f64`, or transfers without copying when the existing stride is +//! already compatible. +//! - Two-dimensional `NdArray` to [`Table`]: copies each `f64` column into a +//! separate 64-byte-aligned `FloatArray`. +//! - One-dimensional `NdArray` to [`Array`]: moves the `f64` buffer into a +//! `FloatArray`. +//! - [`Table`] to `NdArray` through [`TryFrom`]: copies `f64` values and converts +//! nulls to `NaN`. +//! - DLPack import and export: shares data with compatible `f32` and `f64` +//! consumers. Whether a transfer is zero-copy depends on ownership, alignment, +//! layout, and protocol constraints. + +use std::fmt; +use std::ops::{Index, IndexMut, Range, RangeFrom, RangeFull, RangeTo}; +use std::sync::Arc; + +use crate::enums::error::MinarrowError; +use crate::enums::shape_dim::ShapeDim; +use crate::structs::buffer::Buffer; +#[cfg(all(feature = "views", feature = "select"))] +use crate::traits::selection::{AxisSelection, DataSelector, RowSelection}; +use crate::traits::type_unions::Float; +use crate::traits::{concatenate::Concatenate, shape::Shape}; +use crate::{Array, ArrowType, Field, FieldArray, FloatArray, NumericArray, Table, Vec64}; + +#[cfg(feature = "matrix")] +use crate::structs::matrix::{Matrix, aligned_stride}; +#[cfg(feature = "views")] +use crate::structs::views::ndarray_view::NdArrayV; +#[cfg(feature = "dlpack")] +use crate::ffi::dlpack::{ + export_to_dlpack, export_to_dlpack_versioned, DLPackTensor, DLPackTensorVersioned, +}; + +// **************************************************************** +// NdArray +// **************************************************************** + +/// N-dimensional dense array of float values. +/// +/// Backed by [`Buffer`] for zero-copy interop with external memory. +/// Compact column-major layout with a 64-byte aligned allocation start. +/// +/// See the [module-level documentation](self) for design rationale. +#[derive(Clone)] +pub struct NdArray { + pub(crate) data: Arc>, + pub(crate) dims: NdDims, + pub name: Option, +} + +/// Logical equality over shape and values in logical order. Strides and +/// the name do not affect equality, matching the view and SuperNdArray types. +impl PartialEq for NdArray { + fn eq(&self, other: &Self) -> bool { + if self.dims.shape() != other.dims.shape() { + return false; + } + self.into_iter() + .zip(other.into_iter()) + .all(|(a, b)| a == b) + } +} + +// *** Construction ************************************************ + +impl NdArray { + /// Create a zeroed NdArray with the given shape, column-major strides. + /// `shape == &[]` creates a rank-zero array containing one scalar value; + /// `shape == &[0]` creates an empty rank-one array. + pub fn new(shape: &[usize]) -> Self { + let dims = NdDims::from_shape(shape); + let total = buffer_len(dims.shape(), dims.strides()); + let mut v = Vec64::with_capacity(total); + v.0.resize(total, T::default()); + NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None } + } + + /// Create a zeroed NdArray with a name. + pub fn new_named(shape: &[usize], name: impl Into) -> Self { + let mut arr = Self::new(shape); + arr.name = Some(name.into()); + arr + } + + /// Create from a flat column-major slice and shape. + /// + /// The slice holds `product(shape)` logical elements in column-major + /// order, matching the compact layout, so the data copies straight in. + /// A shape with zero axes (`&[]`) has product one and therefore requires + /// one value. A shape containing a zero-length axis, such as `&[0]`, has + /// product zero and requires no values. + pub fn from_slice(data: &[T], shape: &[usize]) -> Self { + let logical_len: usize = shape.iter().product(); + assert_eq!( + data.len(), logical_len, + "NdArray::from_slice: data length {} does not match shape product {}", + data.len(), logical_len + ); + let dims = NdDims::from_shape(shape); + NdArray { + data: Arc::new(Buffer::from_slice(data)), + dims, + name: None, + } + } + + /// Create from an owned 64-byte aligned vector, moving it in without + /// a copy. The vector holds `product(shape)` logical elements in + /// column-major order, matching the compact layout. A shape with zero + /// axes (`&[]`) requires one value, while `&[0]` requires none. + pub fn from_vec64(data: Vec64, shape: &[usize]) -> Self { + let logical_len: usize = shape.iter().product(); + assert_eq!( + data.len(), logical_len, + "NdArray::from_vec64: data length {} does not match shape product {}", + data.len(), logical_len + ); + let dims = NdDims::from_shape(shape); + NdArray { + data: Arc::new(Buffer::from_vec64(data)), + dims, + name: None, + } + } + + /// Create from a pre-owned `Buffer` with explicit shape and strides. + /// + /// The buffer must already contain `buffer_len(shape, strides)` elements + /// in the correct strided layout. + pub fn from_buffer(data: Buffer, shape: &[usize], strides: &[usize]) -> Self { + let required = buffer_len(shape, strides); + assert!( + data.len() >= required, + "NdArray::from_buffer: buffer has {} elements but shape requires {}", + data.len(), required + ); + let dims = NdDims::from_shape_and_strides(shape, strides); + NdArray { data: Arc::new(data), dims, name: None } + } + + /// Create an NdArray filled with a constant value. + pub fn fill(shape: &[usize], value: T) -> Self { + let dims = NdDims::from_shape(shape); + let total = buffer_len(dims.shape(), dims.strides()); + let mut v = Vec64::with_capacity(total); + v.0.resize(total, value); + NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None } + } + + /// Create an NdArray of ones. + pub fn ones(shape: &[usize]) -> Self { + Self::fill(shape, T::one()) + } + + /// Create a 2D identity matrix. + pub fn eye(n: usize) -> Self { + let mut arr = Self::new(&[n, n]); + let stride = arr.dims.strides()[1]; + let buf = Arc::make_mut(&mut arr.data).as_mut_slice(); + for i in 0..n { + buf[i * stride + i] = T::one(); + } + arr + } + + /// Create a 1D NdArray with evenly spaced values in `[start, end]`. + pub fn linspace(start: T, end: T, n: usize) -> Self { + assert!(n >= 2, "linspace requires at least 2 points"); + let step = (end - start) / T::from(n - 1).unwrap(); + let v: Vec64 = (0..n).map(|i| start + step * T::from(i).unwrap()).collect(); + NdArray { + data: Arc::new(Buffer::from_vec64(v)), + dims: NdDims::from_shape(&[n]), + name: None, + } + } + + /// Create a 1D NdArray with values `start, start+step, start+2*step, ...` + /// for `n` elements. + pub fn arange(start: T, step: T, n: usize) -> Self { + let v: Vec64 = (0..n).map(|i| start + step * T::from(i).unwrap()).collect(); + NdArray { + data: Arc::new(Buffer::from_vec64(v)), + dims: NdDims::from_shape(&[n]), + name: None, + } + } + + // *** Shape and introspection ********************************* + + /// Number of dimensions. + #[inline] + pub fn ndim(&self) -> usize { self.dims.ndim() } + + /// Shape as a slice of dimension sizes. + #[inline] + pub fn shape(&self) -> &[usize] { self.dims.shape() } + + /// Strides as a slice of element offsets per dimension. + #[inline] + pub fn strides(&self) -> &[usize] { self.dims.strides() } + + /// Total number of logical elements i.e. the product of shape. + #[inline] + pub fn len(&self) -> usize { self.dims.len() } + + /// Leading-axis (axis 0) observation count i.e. `shape()[0]`, + /// matching `SuperNdArray::n_obs` and NumPy's `len(arr)`. + /// Panics for a rank-zero scalar, which has no axis 0. + #[inline] + pub fn n_obs(&self) -> usize { + assert!(self.ndim() > 0, "n_obs() requires an axis 0"); + self.dims.shape()[0] + } + + /// True if any dimension is zero. + #[inline] + pub fn is_empty(&self) -> bool { self.len() == 0 } + + /// True if the buffer layout matches compact column-major strides, + /// with no transposition or non-standard stride pattern. + pub fn is_contiguous(&self) -> bool { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + let mut expected = 1; + for d in 0..shape.len() { + if strides[d] != expected { + return false; + } + expected *= shape[d]; + } + true + } + + /// True if any element is NaN. + pub fn has_nan(&self) -> bool { + self.into_iter().any(|v| v.is_nan()) + } + + /// Count of NaN elements. + pub fn nan_count(&self) -> usize { + self.into_iter().filter(|v| v.is_nan()).count() + } + + // *** Element access ****************************************** + + /// Get element by N-dimensional index. Panics if out of bounds. + #[inline] + pub fn get(&self, indices: &[usize]) -> T { + self.data.as_slice()[self.offset_of(indices)] + } + + /// Like `get`, but skips bounds checks. + /// + /// # Safety + /// The caller guarantees each index is within its dimension. + #[inline(always)] + pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T { + let strides = self.dims.strides(); + let mut off = 0; + for d in 0..indices.len() { + off += indices[d] * strides[d]; + } + // SAFETY: in-bounds indices produce an in-bounds flat offset. + unsafe { *self.data.as_slice().get_unchecked(off) } + } + + /// Set element by N-dimensional index. Panics if out of bounds. + /// Triggers copy-on-write when views share the buffer. + /// + /// For repeated writes, detach once and write through the slice: + /// ```ignore + /// let s1 = a.strides()[1]; + /// let s = a.as_mut_slice(); // one copy-on-write detach + /// for (i, j, v) in writes { + /// // SAFETY: i and j are within shape + /// unsafe { *s.get_unchecked_mut(i + j * s1) = v; } + /// } + /// ``` + #[inline] + pub fn set(&mut self, indices: &[usize], value: T) { + let off = self.offset_of(indices); + Arc::make_mut(&mut self.data).as_mut_slice()[off] = value; + } + + /// Compute flat buffer offset for an N-dimensional index. + #[inline] + pub(crate) fn offset_of(&self, indices: &[usize]) -> usize { + offset_of_impl(indices, self.dims.shape(), self.dims.strides()) + } + + // *** Metadata ************************************************ + + /// Set the array name. + #[inline] + pub fn set_name(&mut self, name: impl Into) { + self.name = Some(name.into()); + } + + /// Immutable reference to the full flat buffer. + #[inline] + pub fn as_slice(&self) -> &[T] { + self.data.as_slice() + } + + /// Mutable reference to the full flat buffer. Triggers copy-on-write + /// when views share the buffer. + #[inline] + pub fn as_mut_slice(&mut self) -> &mut [T] { + Arc::make_mut(&mut self.data).as_mut_slice() + } + + /// Fill every logical element with a value. + pub fn fill_with(&mut self, value: T) { + // For contiguous arrays, fill the whole buffer + if self.is_contiguous() { + Arc::make_mut(&mut self.data).as_mut_slice().fill(value); + return; + } + // Walk logical positions for non-contiguous layouts + let offsets: Vec = { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + let mut result = Vec::with_capacity(self.len()); + let mut indices = vec![0usize; shape.len()]; + for _ in 0..self.len() { + result.push(indices.iter().zip(strides).map(|(&i, &s)| i * s).sum()); + let mut carry = true; + for d in 0..shape.len() { + if carry { + indices[d] += 1; + if indices[d] < shape[d] { carry = false; } else { indices[d] = 0; } + } + } + } + result + }; + let buf = Arc::make_mut(&mut self.data).as_mut_slice(); + for off in offsets { buf[off] = value; } + } + + // *** 2D axis access (column-major) *************************** + + /// Immutable column slice for a 2D array. + /// Panics if ndim != 2, the column index is out of bounds, or axis 0 + /// is not unit-stride i.e. column elements are not contiguous. + #[inline] + pub fn col(&self, col: usize) -> &[T] { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "col() requires a 2D array"); + assert!(col < shape[1], "Column index out of bounds"); + assert_eq!( + self.dims.strides()[0], 1, + "col() requires unit stride on axis 0; call to_contiguous() first" + ); + let stride = self.dims.strides()[1]; + let start = col * stride; + &self.data.as_slice()[start..start + shape[0]] + } + + /// Mutable column slice for a 2D array. + /// Panics if ndim != 2, the column index is out of bounds, or axis 0 + /// is not unit-stride i.e. column elements are not contiguous. + /// Triggers copy-on-write if the buffer is shared. + #[inline] + pub fn col_mut(&mut self, col: usize) -> &mut [T] { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "col_mut() requires a 2D array"); + assert!(col < shape[1], "Column index out of bounds"); + assert_eq!( + self.dims.strides()[0], 1, + "col_mut() requires unit stride on axis 0; call to_contiguous() first" + ); + let stride = self.dims.strides()[1]; + let n_rows = shape[0]; + let start = col * stride; + &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows] + } + + /// All columns as immutable slices. 2D only. + /// Panics if ndim != 2 or axis 0 is not unit-stride. + pub fn columns(&self) -> Vec<&[T]> { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "columns() requires a 2D array"); + assert_eq!( + self.dims.strides()[0], 1, + "columns() requires unit stride on axis 0; call to_contiguous() first" + ); + let stride = self.dims.strides()[1]; + let n_rows = shape[0]; + let buf = self.data.as_slice(); + (0..shape[1]) + .map(|c| &buf[c * stride..c * stride + n_rows]) + .collect() + } + + /// All columns as mutable slices. 2D only. + /// Panics if ndim != 2 or axis 0 is not unit-stride. + pub fn columns_mut(&mut self) -> Vec<&mut [T]> { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "columns_mut() requires a 2D array"); + assert_eq!( + self.dims.strides()[0], 1, + "columns_mut() requires unit stride on axis 0; call to_contiguous() first" + ); + let n_rows = shape[0]; + let n_cols = shape[1]; + let stride = self.dims.strides()[1]; + let ptr = Arc::make_mut(&mut self.data).as_mut_slice().as_mut_ptr(); + let mut result = Vec::with_capacity(n_cols); + for c in 0..n_cols { + let start = c * stride; + // SAFETY: each slice is within bounds and non-overlapping, + // we have exclusive &mut access. + unsafe { + let col_ptr = ptr.add(start); + result.push(std::slice::from_raw_parts_mut(col_ptr, n_rows)); + } + } + result + } + + /// Wrap as a full `NdArrayV` view for repeated observation access. + /// + /// Call `.obs(i)` on the returned view to get individual observations. + /// The view holds the parent through the array's shared internal + /// buffer, so this is a refcount bump. + #[cfg(feature = "views")] + pub fn as_view(&self) -> NdArrayV { + NdArrayV::from_ndarray(self.clone()) + } + + /// Zero-copy view of a single observation (axis-0 element). + /// + /// Returns an (N-1)-dimensional `NdArrayV` view. For a 2D array + /// with shape `[n, m]`, returns a 1D view of shape `[m]`. For 3D + /// `[n, m, k]`, returns 2D `[m, k]`. Requires rank 2 or higher - for + /// scalar access on a 1D array use `get(&[i])`. + /// + /// For repeated access in a loop, prefer `nd.as_view()` then call + /// `.obs()` on the view to avoid re-wrapping each time. + #[cfg(feature = "views")] + pub fn obs(&self, idx: usize) -> NdArrayV { + self.as_view().obs(idx) + } + + // *** BLAS/LAPACK compatibility (2D) ************************** + + /// Number of rows as i32. Panics if ndim != 2. + #[inline] + pub fn m(&self) -> i32 { + assert_eq!(self.ndim(), 2, "m() requires a 2D array"); + self.dims.shape()[0] as i32 + } + + /// Number of columns as i32. Panics if ndim != 2. + #[inline] + pub fn n(&self) -> i32 { + assert_eq!(self.ndim(), 2, "n() requires a 2D array"); + self.dims.shape()[1] as i32 + } + + /// Leading dimension for BLAS. Panics if ndim != 2. + #[inline] + pub fn lda(&self) -> i32 { + assert_eq!(self.ndim(), 2, "lda() requires a 2D array"); + self.dims.strides()[1] as i32 + } + + // *** Reshape ************************************************* + + /// Reshape to a new shape. Returns a new NdArray with re-laid out data. + /// + /// The total number of logical elements must match. Data is copied + /// from logical element order into the new strided layout. + pub fn reshape(&self, new_shape: &[usize]) -> Result, MinarrowError> { + let new_len: usize = new_shape.iter().product(); + if new_len != self.len() { + return Err(MinarrowError::ShapeError { + message: format!( + "reshape: cannot reshape array of size {} into shape {:?}", + self.len(), new_shape + ), + }); + } + let flat: Vec64 = self.into_iter().collect(); + let mut result = NdArray::from_slice(&flat, new_shape); + result.name = self.name.clone(); + Ok(result) + } + + /// Transpose by reversing the axis order. Returns a new NdArray with + /// re-laid out data. A 1D array copies through unchanged. + /// + /// For a zero-copy transposed view, call `as_view()` and use the + /// view's `transpose()`. + pub fn transpose(&self) -> NdArray { + let shape = self.dims.shape(); + let ndim = shape.len(); + if ndim <= 1 { + let mut result = self.to_contiguous(); + result.name = self.name.clone(); + return result; + } + if ndim == 2 { + let (n_rows, n_cols) = (shape[0], shape[1]); + let new_shape = [n_cols, n_rows]; + let mut result = NdArray::new(&new_shape); + let src_stride = self.dims.strides()[1]; + let dst_stride = result.dims.strides()[1]; + let src = self.data.as_slice(); + let dst = Arc::make_mut(&mut result.data).as_mut_slice(); + for r in 0..n_rows { + for c in 0..n_cols { + dst[r * dst_stride + c] = src[c * src_stride + r]; + } + } + result.name = self.name.clone(); + return result; + } + + // General N-D. Walking the result's logical positions in column-major + // order reads the source at the reversed index, which is the source + // walked with reversed strides. + let new_shape: Vec = shape.iter().rev().copied().collect(); + let rev_strides: Vec = self.dims.strides().iter().rev().copied().collect(); + let src = self.data.as_slice(); + let total = self.len(); + let mut buf = Vec64::with_capacity(total); + let mut indices = vec![0usize; ndim]; + for _ in 0..total { + let offset: usize = indices.iter() + .zip(rev_strides.iter()) + .map(|(&i, &s)| i * s) + .sum(); + buf.push(src[offset]); + let mut carry = true; + for d in 0..ndim { + if carry { + indices[d] += 1; + if indices[d] < new_shape[d] { + carry = false; + } else { + indices[d] = 0; + } + } + } + } + NdArray { + data: Arc::new(Buffer::from_vec64(buf)), + dims: NdDims::from_shape(&new_shape), + name: self.name.clone(), + } + } + + /// Flatten to a contiguous 1D array. + pub fn flatten(&self) -> NdArray { + let flat: Vec64 = self.into_iter().collect(); + let n = flat.len(); + NdArray { + data: Arc::new(Buffer::from_vec64(flat)), + dims: NdDims::from_shape(&[n]), + name: self.name.clone(), + } + } + + /// If the array has non-standard strides, re-lay out into default + /// column-major contiguous form. + pub fn to_contiguous(&self) -> NdArray { + if self.is_contiguous() { + return self.clone(); + } + let mut result = NdArray::from_slice( + &self.into_iter().collect::>(), + self.shape(), + ); + result.name = self.name.clone(); + result + } + + // *** Slicing: arr.slice(nd![1..4, 2..5]) ************************* + + /// Slice this array along any combination of axes. + /// + /// Each axis takes any [`DataSelector`] - a single index collapses + /// that dimension, and a contiguous range keeps it. Returns a + /// zero-copy `NdArrayV` view that holds the parent through the + /// array's shared internal buffer, so this is a refcount bump, + /// matching `as_view` and `obs`. + /// + /// # Examples + /// ```ignore + /// arr.slice(&[&2]) // single index on 1D + /// arr.slice(&[&(1..4)]) // range on axis 0 + /// arr.slice(nd![1..4, 2..5]) // range on both axes (2D) + /// arr.slice(nd![0..3, 2]) // range on axis 0, single on axis 1 + /// arr.slice(nd![1, 0..4, 3]) // mixed for 3D + /// ``` + #[cfg(all(feature = "views", feature = "select"))] + pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + assert_eq!( + selection.len(), self.ndim(), + "slice(): expected {} axes, got {}", self.ndim(), selection.len() + ); + + let shape = self.dims.shape(); + let strides = self.dims.strides(); + + // Compute the new offset, shape, and strides + let mut new_offset: usize = 0; + let mut new_shape = Vec::with_capacity(self.ndim()); + let mut new_strides = Vec::with_capacity(self.ndim()); + + for (d, sel) in selection.iter().enumerate() { + let (start, end, collapse) = sel.resolve_axis(shape[d]); + assert!( + end <= shape[d], + "slice(): end {} out of bounds for axis {} (size {})", end, d, shape[d] + ); + new_offset += start * strides[d]; + if !collapse { + new_shape.push(end - start); + new_strides.push(strides[d]); + } + } + + NdArrayV::new(self.clone(), new_offset, &new_shape, &new_strides) + } + + // *** Apply *************************************************** + + /// Apply a function to every logical element, returning a new compact + /// array with this array's shape and name. The caller supplies the + /// operation; NdArray supplies logical-order traversal and materialisation. + pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray { + let flat: Vec64 = self.into_iter().map(f).collect(); + let mut result = NdArray::from_slice(&flat, self.shape()); + result.name = self.name.clone(); + result + } + + /// Apply a function to every logical element in place, with no + /// allocation. Copy-on-write triggers first when views share the + /// buffer. + pub fn apply_mut(&mut self, f: impl Fn(T) -> T) { + if self.is_contiguous() { + for v in Arc::make_mut(&mut self.data).as_mut_slice() { + *v = f(*v); + } + return; + } + // Walk logical positions for non-contiguous layouts so stride + // padding stays untouched. + let shape = self.dims.shape().to_vec(); + let strides = self.dims.strides().to_vec(); + let total = self.len(); + let buf = Arc::make_mut(&mut self.data).as_mut_slice(); + let mut indices = vec![0usize; shape.len()]; + for _ in 0..total { + let offset: usize = indices.iter().zip(strides.iter()).map(|(&i, &s)| i * s).sum(); + buf[offset] = f(buf[offset]); + let mut carry = true; + for d in 0..shape.len() { + if carry { + indices[d] += 1; + if indices[d] < shape[d] { + carry = false; + } else { + indices[d] = 0; + } + } + } + } + } + + /// Apply a function to every 1D lane along the given axis, collapsing + /// that axis. Each lane arrives as a zero-copy [`NdArrayV`] and the + /// closure returns one value for it, so the output shape drops `axis`. + /// Requires rank 2 or higher - a 1D array is itself a single lane. + #[cfg(feature = "views")] + pub fn apply_axis(&self, axis: usize, mut f: impl FnMut(NdArrayV) -> T) -> NdArray { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + let ndim = shape.len(); + assert!(ndim >= 2, "apply_axis requires a 2D or higher array"); + assert!(axis < ndim, "apply_axis: axis {} out of bounds for {}D array", axis, ndim); + + let out_shape: Vec = shape + .iter() + .enumerate() + .filter(|(d, _)| *d != axis) + .map(|(_, &s)| s) + .collect(); + let out_dims: Vec = (0..ndim).filter(|&d| d != axis).collect(); + let total: usize = out_shape.iter().product(); + + let lane_shape = [shape[axis]]; + let lane_strides = [strides[axis]]; + + // Walk the output positions in column-major order. Each position + // anchors one lane's base offset in the source. + let mut flat = Vec64::with_capacity(total); + let mut indices = vec![0usize; out_shape.len()]; + for _ in 0..total { + let offset: usize = indices + .iter() + .zip(out_dims.iter()) + .map(|(&i, &d)| i * strides[d]) + .sum(); + let lane = NdArrayV::new(self.clone(), offset, &lane_shape, &lane_strides); + flat.push(f(lane)); + let mut carry = true; + for d in 0..out_shape.len() { + if carry { + indices[d] += 1; + if indices[d] < out_shape[d] { + carry = false; + } else { + indices[d] = 0; + } + } + } + } + let mut result = NdArray::from_slice(&flat, &out_shape); + result.name = self.name.clone(); + result + } + + // *** Conversions ********************************************* + + /// Export as a legacy DLPack tensor for PyTorch, NumPy, JAX, and other + /// compatible consumers. Shared or aliased storage copies because the + /// legacy protocol cannot mark the exported tensor read-only. + /// + /// Returns a `DLPackTensor` that manages the lifecycle. Drop it to + /// release, or call `.into_raw()` to transfer ownership to an FFI + /// consumer such as a PyCapsule. + #[cfg(feature = "dlpack")] + pub fn to_dlpack(self) -> DLPackTensor { + export_to_dlpack(self) + } + + /// Export as a DLPack 1.x versioned tensor, carrying the spec version + /// and flags fields that PyTorch and JAX negotiate. The read-only + /// flag is set when the backing buffer is shared. + #[cfg(feature = "dlpack")] + pub fn to_dlpack_versioned(self) -> DLPackTensorVersioned { + export_to_dlpack_versioned(self) + } + + // *** Parallel iteration (rayon) ****************************** + + /// Parallel iterator over the underlying buffer. Rayon splits + /// the contiguous data into chunks across threads automatically. + #[cfg(feature = "parallel_proc")] + pub fn par_iter(&self) -> rayon::slice::Iter<'_, T> + where + T: Send + Sync, + { + use rayon::prelude::*; + self.data.as_slice().par_iter() + } + + /// Parallel chunks of the underlying buffer. Each chunk is a + /// contiguous `&[T]` slice that rayon distributes across threads. + #[cfg(feature = "parallel_proc")] + pub fn par_chunks(&self, chunk_size: usize) -> rayon::slice::Chunks<'_, T> + where + T: Send + Sync, + { + use rayon::prelude::*; + self.data.as_slice().par_chunks(chunk_size) + } + + /// Iterator over axis-0 observations. Each item is the observation + /// index and a zero-copy `NdArrayV` view. + #[cfg(feature = "views")] + pub fn iter_obs(&self) -> impl Iterator)> + '_ { + assert!(self.ndim() >= 2, "iter_obs() requires a 2D or higher array"); + let n_obs = self.shape()[0]; + (0..n_obs).map(move |i| (i, self.obs(i))) + } + + /// Parallel iterator over axis-0 observations. Each item is the + /// observation index and a zero-copy `NdArrayV` view. + #[cfg(all(feature = "parallel_proc", feature = "views"))] + pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator)> + '_ + where + T: Send + Sync, + { + use rayon::prelude::*; + assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher array"); + let n_obs = self.dims.shape()[0]; + (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i))) + } + + /// Iterate one logical axis-0 run identified by its flattened outer + /// index. This composes the logical iterator for SuperNdArray without + /// materialising its batches. + pub(crate) fn iter_axis0_run(&self, run_idx: usize) -> impl ExactSizeIterator + '_ { + assert!(self.ndim() > 0, "axis-0 iteration requires an axis 0"); + let n_runs: usize = self.shape()[1..].iter().product(); + assert!(run_idx < n_runs, "axis-0 run {} out of bounds ({})", run_idx, n_runs); + + let mut rem = run_idx; + let mut offset = 0; + for d in 1..self.ndim() { + offset += (rem % self.shape()[d]) * self.strides()[d]; + rem /= self.shape()[d]; + } + let stride = self.strides()[0]; + (0..self.shape()[0]).map(move |i| self.data.as_slice()[offset + i * stride]) + } +} + +// *** f64 conversions to the Array/Table/Matrix enum boundary ***** + +impl NdArray { + /// Convert a 2D NdArray to a Table. + /// + /// Each column is copied into its own 64-byte aligned `FloatArray`, + /// since the compact tensor layout does not place column starts on + /// alignment boundaries. `fields` must have exactly `n_cols` entries + /// when supplied. Each supplied field must describe a non-nullable + /// [`ArrowType::Float64`] column. `None` generates `col_0`, `col_1`, + /// and so on. + pub fn to_table(self, fields: Option>) -> Result { + let shape = self.dims.shape(); + if shape.len() != 2 { + return Err(MinarrowError::ShapeError { + message: format!("to_table requires a 2D array, got {}D", shape.len()), + }); + } + let n_cols = shape[1]; + let fields = match fields { + Some(fields) => { + if fields.len() != n_cols { + return Err(MinarrowError::ShapeError { + message: format!( + "to_table: expected {} fields for {} columns, got {}", + n_cols, n_cols, fields.len() + ), + }); + } + for (index, field) in fields.iter().enumerate() { + if field.dtype != ArrowType::Float64 { + return Err(MinarrowError::TypeError { + from: "Field", + to: "Float64 NdArray column", + message: Some(format!( + "to_table: field {} ('{}') has dtype {:?}, expected Float64", + index, field.name, field.dtype + )), + }); + } + if field.nullable { + return Err(MinarrowError::NullError { + message: Some(format!( + "to_table: field {} ('{}') is nullable, but NdArray columns are non-nullable", + index, field.name + )), + }); + } + } + fields + } + None => (0..n_cols) + .map(|i| Field::new(format!("col_{}", i), ArrowType::Float64, false, None)) + .collect(), + }; + // The column copies below read unit-stride axis-0 runs, so a + // transposed or otherwise strided layout re-lays first. + let this = if self.dims.strides()[0] == 1 { self } else { self.to_contiguous() }; + let n_rows = this.dims.shape()[0]; + let stride = this.dims.strides()[1]; + let name = this.name; + let buf = this.data.as_slice(); + + let mut cols = Vec::with_capacity(n_cols); + for (i, field) in fields.into_iter().enumerate() { + let col_start = i * stride; + let col: Buffer = Buffer::from_slice(&buf[col_start..col_start + n_rows]); + let float_arr = FloatArray::new(col, None); + let array = Array::NumericArray(NumericArray::Float64(Arc::new(float_arr))); + cols.push(FieldArray::new(field, array)); + } + + Ok(Table::new(name.unwrap_or_default(), Some(cols))) + } + + /// Convert a 1D NdArray to an Array (FloatArray). + /// Compact, exact-length storage moves across directly. A strided array + /// or one with unused backing elements is materialised in logical order. + pub fn to_array(self) -> Result { + if self.ndim() != 1 { + return Err(MinarrowError::ShapeError { + message: format!("to_array requires a 1D array, got {}D", self.ndim()), + }); + } + let buffer = if self.dims.strides()[0] == 1 && self.data.len() == self.len() { + Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone()) + } else { + Buffer::from_vec64((&self).into_iter().collect()) + }; + let float_arr = FloatArray::new(buffer, None); + Ok(Array::NumericArray(NumericArray::Float64(Arc::new(float_arr)))) + } + + /// Convert a 2D NdArray to a Matrix. + /// + /// Matrix pads each column to a 64-byte boundary for BLAS/LAPACK and + /// SIMD access, so compact tensor data is re-laid out into the padded + /// form. An array whose stride already matches the padded layout, such + /// as one built from a Matrix, moves across zero-copy. + #[cfg(feature = "matrix")] + pub fn to_matrix(self) -> Result { + let shape = self.dims.shape(); + if shape.len() != 2 { + return Err(MinarrowError::ShapeError { + message: format!("to_matrix requires a 2D array, got {}D", shape.len()), + }); + } + let n_rows = shape[0]; + let n_cols = shape[1]; + let strides = self.dims.strides(); + + if strides[0] == 1 && strides[1] == aligned_stride(n_rows) { + return Ok(Matrix { + n_rows, + n_cols, + stride: strides[1], + data: Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone()), + name: self.name, + }); + } + + let name = self.name.clone(); + let compact: Vec64 = self.into_iter().collect(); + Ok(Matrix::from_f64_unaligned(&compact, n_rows, n_cols, name)) + } +} + +// **************************************************************** +// NdDims - internal dimension storage +// **************************************************************** + +/// Internal storage for shape and strides. +/// +/// Inline arrays for 1D-5D to avoid heap allocation; the `Dn` variant +/// handles 6+ dimensions via boxed slices. +#[derive(Clone, PartialEq)] +pub(crate) enum NdDims { + D1 { shape: [usize; 1], strides: [usize; 1] }, + D2 { shape: [usize; 2], strides: [usize; 2] }, + D3 { shape: [usize; 3], strides: [usize; 3] }, + D4 { shape: [usize; 4], strides: [usize; 4] }, + D5 { shape: [usize; 5], strides: [usize; 5] }, + Dn { shape: Box<[usize]>, strides: Box<[usize]> }, +} + +impl NdDims { + /// Build dims from a shape slice, computing compact column-major + /// strides. + pub(crate) fn from_shape(shape: &[usize]) -> Self { + let strides = col_major_strides(shape); + Self::from_shape_and_strides(shape, &strides) + } + + /// Build dims from explicit shape and strides. + pub(crate) fn from_shape_and_strides(shape: &[usize], strides: &[usize]) -> Self { + assert_eq!( + shape.len(), + strides.len(), + "NdArray: shape rank {} does not match strides rank {}", + shape.len(), + strides.len() + ); + match shape.len() { + 1 => NdDims::D1 { + shape: [shape[0]], + strides: [strides[0]], + }, + 2 => NdDims::D2 { + shape: [shape[0], shape[1]], + strides: [strides[0], strides[1]], + }, + 3 => NdDims::D3 { + shape: [shape[0], shape[1], shape[2]], + strides: [strides[0], strides[1], strides[2]], + }, + 4 => NdDims::D4 { + shape: [shape[0], shape[1], shape[2], shape[3]], + strides: [strides[0], strides[1], strides[2], strides[3]], + }, + 5 => NdDims::D5 { + shape: [shape[0], shape[1], shape[2], shape[3], shape[4]], + strides: [strides[0], strides[1], strides[2], strides[3], strides[4]], + }, + _ => NdDims::Dn { + shape: shape.into(), + strides: strides.into(), + }, + } + } + + /// Number of dimensions. + #[inline] + pub(crate) fn ndim(&self) -> usize { + match self { + NdDims::D1 { .. } => 1, + NdDims::D2 { .. } => 2, + NdDims::D3 { .. } => 3, + NdDims::D4 { .. } => 4, + NdDims::D5 { .. } => 5, + NdDims::Dn { shape, .. } => shape.len(), + } + } + + /// Shape as a slice. + #[inline] + pub(crate) fn shape(&self) -> &[usize] { + match self { + NdDims::D1 { shape, .. } => shape, + NdDims::D2 { shape, .. } => shape, + NdDims::D3 { shape, .. } => shape, + NdDims::D4 { shape, .. } => shape, + NdDims::D5 { shape, .. } => shape, + NdDims::Dn { shape, .. } => shape, + } + } + + /// Strides as a slice. + #[inline] + pub(crate) fn strides(&self) -> &[usize] { + match self { + NdDims::D1 { strides, .. } => strides, + NdDims::D2 { strides, .. } => strides, + NdDims::D3 { strides, .. } => strides, + NdDims::D4 { strides, .. } => strides, + NdDims::D5 { strides, .. } => strides, + NdDims::Dn { strides, .. } => strides, + } + } + + /// Total logical element count i.e. the product of all dimensions. + #[inline] + pub(crate) fn len(&self) -> usize { + self.shape().iter().product() + } +} + +// **************************************************************** +// Axis selection input +// **************************************************************** + +/// Build an axis-selection slice from mixed indices and ranges. +/// +/// Each entry is any [`DataSelector`](crate::traits::selection::DataSelector) - +/// a single index collapses the dimension, and a contiguous range keeps it. +/// +/// # Example +/// ```ignore +/// arr.slice(nd![0..3, 2, 1..4]) +/// ``` +#[macro_export] +macro_rules! nd { + ($($sel:expr),+ $(,)?) => { + &[$(&$sel as &dyn $crate::traits::selection::DataSelector),+] + }; +} + +// **************************************************************** +// Stride computation +// **************************************************************** + +/// Compute compact column-major strides with no inter-dimension padding. +/// +/// For a shape `[a, b, c]`, the strides are `[1, a, a * b]`. The buffer is +/// fully contiguous, so DLPack consumers receive a dense tensor and range +/// indexing over the outermost axis reads logical data with no gaps. The +/// backing allocation start remains 64-byte aligned through `Vec64`. +pub(crate) fn col_major_strides(shape: &[usize]) -> Vec { + let mut strides = Vec::with_capacity(shape.len()); + if shape.is_empty() { + return strides; + } + strides.push(1); + for d in 1..shape.len() { + strides.push(strides[d - 1] * shape[d - 1]); + } + strides +} + +/// Total buffer length required for a given shape and strides. +pub(crate) fn buffer_len(shape: &[usize], strides: &[usize]) -> usize { + if shape.iter().any(|&d| d == 0) { + return 0; + } + // The last element is at sum((shape[d]-1) * strides[d]) for all d. + // Buffer must hold one past that. + let max_offset: usize = shape.iter() + .zip(strides.iter()) + .map(|(&s, &st)| (s - 1) * st) + .sum(); + max_offset + 1 +} + +// **************************************************************** +// IntoIterator +// **************************************************************** + +/// Iterating an NdArray yields `T` values in column-major order, +/// walking contiguous runs along axis 0 (the innermost dimension) +/// and advancing through higher dimensions. Each column/slice is +/// a sequential cache-friendly read with no per-element arithmetic. +impl<'a, T: Float> IntoIterator for &'a NdArray { + type Item = T; + type IntoIter = NdArrayIter<'a, T>; + + #[inline] + fn into_iter(self) -> NdArrayIter<'a, T> { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + if shape.is_empty() { + return NdArrayIter { + buf: self.data.as_slice(), + n_inner: 1, + inner_stride: 1, + run_offsets: vec![0], + run_idx: 0, + inner_idx: 0, + total: 1, + yielded: 0, + }; + } + let n_inner = shape[0]; + + // Number of contiguous runs = product of all dims except axis 0 + let n_runs: usize = shape[1..].iter().product(); + + // Build the starting offset of each contiguous run. + // For 1D there is one run at offset 0. + // For 2D these are just [0, stride1, 2*stride1, ...]. + // For N-D we walk the outer indices in column-major order. + let mut run_offsets = Vec::with_capacity(n_runs); + if shape.len() <= 1 { + run_offsets.push(0); + } else { + let outer_shape = &shape[1..]; + let outer_strides = &strides[1..]; + let mut outer_indices = vec![0usize; outer_shape.len()]; + for _ in 0..n_runs { + let off: usize = outer_indices.iter() + .zip(outer_strides.iter()) + .map(|(&i, &s)| i * s) + .sum(); + run_offsets.push(off); + // Advance outer indices (column-major) + let mut carry = true; + for d in 0..outer_shape.len() { + if carry { + outer_indices[d] += 1; + if outer_indices[d] < outer_shape[d] { + carry = false; + } else { + outer_indices[d] = 0; + } + } + } + } + } + + NdArrayIter { + buf: self.data.as_slice(), + n_inner, + inner_stride: strides[0], + run_offsets, + run_idx: 0, + inner_idx: 0, + total: self.len(), + yielded: 0, + } + } +} + +/// Consuming iterator - collects logical elements then iterates. +impl IntoIterator for NdArray { + type Item = T; + type IntoIter = std::vec::IntoIter; + + #[inline] + fn into_iter(self) -> std::vec::IntoIter { + let v: Vec = (&self).into_iter().collect(); + v.into_iter() + } +} + +/// Iterator over NdArray elements in column-major order. +/// +/// When `inner_stride` is 1 (the normal case), walks contiguous runs +/// along axis 0 with sequential memory reads. When `inner_stride` > 1 +/// (e.g. after collapsing axis 0 via slicing), steps through strided +/// elements within each run. +pub struct NdArrayIter<'a, T> { + pub(crate) buf: &'a [T], + pub(crate) n_inner: usize, + pub(crate) inner_stride: usize, + pub(crate) run_offsets: Vec, + pub(crate) run_idx: usize, + pub(crate) inner_idx: usize, + pub(crate) total: usize, + pub(crate) yielded: usize, +} + +impl<'a, T: Float> Iterator for NdArrayIter<'a, T> { + type Item = T; + + #[inline] + fn next(&mut self) -> Option { + if self.yielded >= self.total { return None; } + + let val = self.buf[self.run_offsets[self.run_idx] + self.inner_idx * self.inner_stride]; + self.yielded += 1; + self.inner_idx += 1; + if self.inner_idx >= self.n_inner { + self.inner_idx = 0; + self.run_idx += 1; + } + Some(val) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let r = self.total - self.yielded; + (r, Some(r)) + } +} + +impl<'a, T: Float> ExactSizeIterator for NdArrayIter<'a, T> {} + +// **************************************************************** +// Internal helpers +// **************************************************************** + +/// Compute flat buffer offset for an N-dimensional index. +/// Panics on rank mismatch or an out-of-bounds index, in release +/// builds included, since a wrong offset can silently read or write +/// another element's slot. +#[inline] +pub(crate) fn offset_of_impl(indices: &[usize], shape: &[usize], strides: &[usize]) -> usize { + assert_eq!( + indices.len(), + shape.len(), + "NdArray: {} indices for a {}D array", + indices.len(), + shape.len() + ); + let mut offset = 0; + for d in 0..shape.len() { + assert!( + indices[d] < shape[d], + "NdArray: index {} out of bounds for dim {} (size {})", + indices[d], d, shape[d] + ); + offset += indices[d] * strides[d]; + } + offset +} + +// **************************************************************** +// Trait implementations +// **************************************************************** + +impl Shape for NdArray { + fn shape(&self) -> ShapeDim { + match self.dims.ndim() { + 0 => ShapeDim::Rank0(1), + 1 => ShapeDim::Rank1(self.dims.shape()[0]), + 2 => ShapeDim::Rank2 { + rows: self.dims.shape()[0], + cols: self.dims.shape()[1], + }, + _ => ShapeDim::RankN(self.dims.shape().to_vec()), + } + } +} + +impl Concatenate for NdArray { + /// Concatenate along axis 0. All other dimensions must match. + fn concat(self, other: Self) -> Result { + let s1 = self.dims.shape(); + let s2 = other.dims.shape(); + if s1.len() != s2.len() { + return Err(MinarrowError::IncompatibleTypeError { + from: "NdArray", + to: "NdArray", + message: Some(format!( + "Cannot concatenate {}D and {}D arrays", s1.len(), s2.len() + )), + }); + } + if s1.is_empty() { + return Err(MinarrowError::ShapeError { + message: "Cannot concatenate rank-zero arrays along axis 0".to_string(), + }); + } + for d in 1..s1.len() { + if s1[d] != s2[d] { + return Err(MinarrowError::IncompatibleTypeError { + from: "NdArray", + to: "NdArray", + message: Some(format!( + "Dimension {} mismatch: {} vs {}", d, s1[d], s2[d] + )), + }); + } + } + + let mut new_shape: Vec = s1.to_vec(); + new_shape[0] += s2[0]; + + // The result name joins both sides, matching Table's concat. + let name = match (&self.name, &other.name) { + (Some(a), Some(b)) => Some(format!("{}+{}", a, b)), + (Some(a), None) => Some(a.clone()), + (None, Some(b)) => Some(b.clone()), + (None, None) => None, + }; + + // 1D: plain append + if s1.len() == 1 { + let mut flat = Vec64::with_capacity(self.len() + other.len()); + flat.extend(&self); + flat.extend(&other); + let mut result = NdArray::from_slice(&flat, &new_shape); + result.name = name; + return Ok(result); + } + + // Fast path for contiguous 2D: interleave columns with memcpy + if s1.len() == 2 && self.dims.strides()[0] == 1 && other.dims.strides()[0] == 1 { + let new_dims = NdDims::from_shape(&new_shape); + let new_stride = new_dims.strides()[1]; + let total = buffer_len(&new_shape, new_dims.strides()); + let mut buf = Vec64::with_capacity(total); + buf.0.resize(total, T::default()); + let dst = buf.as_mut_slice(); + for c in 0..s1[1] { + let dst_start = c * new_stride; + dst[dst_start..dst_start + s1[0]].copy_from_slice(self.col(c)); + dst[dst_start + s1[0]..dst_start + s1[0] + s2[0]].copy_from_slice(other.col(c)); + } + return Ok(NdArray { + data: Arc::new(Buffer::from_vec64(buf)), + dims: new_dims, + name, + }); + } + + // General case: interleave the axis-0 run of each operand per + // outer index, walking both sources in column-major logical order. + let n1 = s1[0]; + let n2 = s2[0]; + let run_len = n1 + n2; + let n_runs: usize = s1[1..].iter().product(); + let new_dims = NdDims::from_shape(&new_shape); + let total = buffer_len(&new_shape, new_dims.strides()); + let mut buf = Vec64::with_capacity(total); + buf.0.resize(total, T::default()); + let dst = buf.as_mut_slice(); + let mut it_a = (&self).into_iter(); + let mut it_b = (&other).into_iter(); + for r in 0..n_runs { + let base = r * run_len; + for i in 0..n1 { + dst[base + i] = it_a.next().unwrap(); + } + for i in 0..n2 { + dst[base + n1 + i] = it_b.next().unwrap(); + } + } + Ok(NdArray { + data: Arc::new(Buffer::from_vec64(buf)), + dims: new_dims, + name, + }) + } +} + +// *** Axis selection: arr.s(nd![1..4, 2]) ************************* + +/// Selection across every axis at once, delegating to `slice`. Single +/// indices collapse their dimension, and contiguous ranges keep it. +/// Zero-copy. +#[cfg(all(feature = "views", feature = "select"))] +impl AxisSelection for NdArray { + type View = NdArrayV; + + fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + self.slice(selection) + } + + fn get_axis_count(&self) -> usize { + self.ndim() + } +} + +// *** Row selection: arr.r(0..10) ********************************* + +/// Axis-0 observation selection. Contiguous ranges return a zero-copy +/// window view. Index arrays gather the selected observations into an +/// owned array wrapped in a full view, matching `Table`'s behaviour. +#[cfg(all(feature = "views", feature = "select"))] +impl RowSelection for NdArray { + type View = NdArrayV; + + fn r(&self, selection: S) -> NdArrayV { + assert!(self.ndim() > 0, "row selection requires an axis 0"); + let n_obs = self.shape()[0]; + let indices = selection.resolve_indices(n_obs); + if selection.is_contiguous() { + let start = indices.first().copied().unwrap_or(0); + let ranges: Vec> = std::iter::once(start..start + indices.len()) + .chain(self.shape()[1..].iter().map(|&n| 0..n)) + .collect(); + let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect(); + return self.slice(&refs); + } + NdArrayV::from_ndarray(gather_obs_impl( + &indices, + self.shape(), + self.name.clone(), + |idx| self.get(idx), + )) + } + + fn get_row_count(&self) -> usize { + self.n_obs() + } +} + +/// Materialise selected axis-0 observations into a compact owned array. +/// Walks the output positions in column-major order, reading each source +/// element through the provided accessor, so any stride layout gathers +/// correctly. +#[cfg(all(feature = "views", feature = "select"))] +pub(crate) fn gather_obs_impl( + indices: &[usize], + shape: &[usize], + name: Option, + get: impl Fn(&[usize]) -> T, +) -> NdArray { + let mut out_shape = shape.to_vec(); + out_shape[0] = indices.len(); + let total: usize = out_shape.iter().product(); + + let mut flat = Vec64::with_capacity(total); + let ndim = shape.len(); + let mut idx = vec![0usize; ndim]; + let inner_runs: usize = shape[1..].iter().product::().max(1); + for _ in 0..inner_runs { + for &obs in indices { + idx[0] = obs; + flat.push(get(&idx)); + } + // Advance the inner multi-index in column-major order. + let mut carry = true; + for d in 1..ndim { + if carry { + idx[d] += 1; + if idx[d] < shape[d] { + carry = false; + } else { + idx[d] = 0; + } + } + } + } + let mut result = NdArray::from_slice(&flat, &out_shape); + result.name = name; + result +} + +// *** Bracket indexing: arr[col][row] ****************************** + +/// `arr[i]` selects along the outermost stored axis. +/// +/// For 1D, returns a single-element slice. +/// For 2D (column-major), `arr[col]` returns the contiguous column +/// as `&[f64]`, so `arr[col][row]` gives `&f64`. +impl Index for NdArray { + type Output = [T]; + + #[inline] + fn index(&self, idx: usize) -> &[T] { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + match shape.len() { + 0 => panic!("NdArray: a rank-zero array has no axis to index with usize"), + 1 => { + assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]); + &self.data.as_slice()[idx..idx + 1] + } + 2 => { + assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]); + let start = idx * strides[1]; + &self.data.as_slice()[start..start + shape[0]] + } + n => { + // Index the outermost axis (last), return the contiguous inner slab + assert!( + self.is_contiguous(), + "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access" + ); + let last = n - 1; + assert!(idx < shape[last], "index out of bounds for axis {}", last); + let start = idx * strides[last]; + &self.data.as_slice()[start..start + strides[last]] + } + } + } +} + +impl IndexMut for NdArray { + #[inline] + fn index_mut(&mut self, idx: usize) -> &mut [T] { + let shape = self.dims.shape().to_vec(); + let strides = self.dims.strides().to_vec(); + match shape.len() { + 0 => panic!("NdArray: a rank-zero array has no axis to index with usize"), + 1 => { + assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[idx..idx + 1] + } + 2 => { + assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]); + let start = idx * strides[1]; + let n_rows = shape[0]; + &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows] + } + n => { + assert!( + self.is_contiguous(), + "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access" + ); + let last = n - 1; + assert!(idx < shape[last], "index out of bounds for axis {}", last); + let start = idx * strides[last]; + &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + strides[last]] + } + } + } +} + +// *** Range indexing: arr[1..4] ************************************ + +/// `arr[start..end]` selects a contiguous range along the outermost axis. +/// +/// For 1D, returns the element slice directly. +/// For 2D and above, selects a range of outermost-axis entries as one +/// contiguous slab of logical data. Requires a contiguous layout, since a +/// padded or transposed stride pattern has no gap-free slab to return. +/// Non-contiguous arrays panic with guidance to use `slice()`. +impl Index> for NdArray { + type Output = [T]; + + #[inline] + fn index(&self, range: Range) -> &[T] { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + match shape.len() { + 0 => panic!("NdArray: a rank-zero array has no axis to range-index"), + 1 => &self.data.as_slice()[range], + _ => { + assert!( + self.is_contiguous(), + "range indexing requires a contiguous layout, use slice() for strided access" + ); + let last = shape.len() - 1; + assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]); + let start = range.start * strides[last]; + let end = range.end * strides[last]; + &self.data.as_slice()[start..end] + } + } + } +} + +impl IndexMut> for NdArray { + #[inline] + fn index_mut(&mut self, range: Range) -> &mut [T] { + let shape = self.dims.shape().to_vec(); + let strides = self.dims.strides().to_vec(); + match shape.len() { + 0 => panic!("NdArray: a rank-zero array has no axis to range-index"), + 1 => &mut Arc::make_mut(&mut self.data).as_mut_slice()[range], + _ => { + assert!( + self.is_contiguous(), + "range indexing requires a contiguous layout, use slice() for strided access" + ); + let last = shape.len() - 1; + assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]); + let start = range.start * strides[last]; + let end = range.end * strides[last]; + &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..end] + } + } + } +} + +impl Index> for NdArray { + type Output = [T]; + + #[inline] + fn index(&self, range: RangeFrom) -> &[T] { + assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index"); + let last_dim = self.dims.shape().len() - 1; + let end = self.dims.shape()[last_dim]; + &self[range.start..end] + } +} + +impl Index> for NdArray { + type Output = [T]; + + #[inline] + fn index(&self, range: RangeTo) -> &[T] { + &self[0..range.end] + } +} + +impl Index for NdArray { + type Output = [T]; + + #[inline] + fn index(&self, _: RangeFull) -> &[T] { + assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index"); + let last_dim = self.dims.shape().len() - 1; + let end = self.dims.shape()[last_dim]; + &self[0..end] + } +} + +// *** Tuple indexing ********************************************** + +impl Index<()> for NdArray { + type Output = T; + #[inline] + fn index(&self, (): ()) -> &T { + &self.data.as_slice()[self.offset_of(&[])] + } +} + +impl Index<(usize,)> for NdArray { + type Output = T; + #[inline] + fn index(&self, (i,): (usize,)) -> &T { + &self.data.as_slice()[self.offset_of(&[i])] + } +} + +impl Index<(usize, usize)> for NdArray { + type Output = T; + #[inline] + fn index(&self, (i, j): (usize, usize)) -> &T { + &self.data.as_slice()[self.offset_of(&[i, j])] + } +} + +impl Index<(usize, usize, usize)> for NdArray { + type Output = T; + #[inline] + fn index(&self, (i, j, k): (usize, usize, usize)) -> &T { + &self.data.as_slice()[self.offset_of(&[i, j, k])] + } +} + +impl Index<(usize, usize, usize, usize)> for NdArray { + type Output = T; + #[inline] + fn index(&self, (i, j, k, l): (usize, usize, usize, usize)) -> &T { + &self.data.as_slice()[self.offset_of(&[i, j, k, l])] + } +} + +impl Index<(usize, usize, usize, usize, usize)> for NdArray { + type Output = T; + #[inline] + fn index(&self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &T { + &self.data.as_slice()[self.offset_of(&[i, j, k, l, m])] + } +} + +impl IndexMut<(usize,)> for NdArray { + #[inline] + fn index_mut(&mut self, (i,): (usize,)) -> &mut T { + let off = self.offset_of(&[i]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +impl IndexMut<()> for NdArray { + #[inline] + fn index_mut(&mut self, (): ()) -> &mut T { + let off = self.offset_of(&[]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +impl IndexMut<(usize, usize)> for NdArray { + #[inline] + fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut T { + let off = self.offset_of(&[i, j]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +impl IndexMut<(usize, usize, usize)> for NdArray { + #[inline] + fn index_mut(&mut self, (i, j, k): (usize, usize, usize)) -> &mut T { + let off = self.offset_of(&[i, j, k]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +impl IndexMut<(usize, usize, usize, usize)> for NdArray { + #[inline] + fn index_mut(&mut self, (i, j, k, l): (usize, usize, usize, usize)) -> &mut T { + let off = self.offset_of(&[i, j, k, l]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +impl IndexMut<(usize, usize, usize, usize, usize)> for NdArray { + #[inline] + fn index_mut(&mut self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &mut T { + let off = self.offset_of(&[i, j, k, l, m]); + &mut Arc::make_mut(&mut self.data).as_mut_slice()[off] + } +} + +// *** From conversions ******************************************** + +/// 1D from a flat slice. +impl From<&[T]> for NdArray { + fn from(data: &[T]) -> Self { + NdArray { + data: Arc::new(Buffer::from_slice(data)), + dims: NdDims::from_shape(&[data.len()]), + name: None, + } + } +} + +/// 1D from owned Vec64. +impl From> for NdArray { + fn from(v: Vec64) -> Self { + let n = v.len(); + NdArray { + data: Arc::new(Buffer::from_vec64(v)), + dims: NdDims::from_shape(&[n]), + name: None, + } + } +} + +/// 2D from column vectors. +impl From<&[Vec]> for NdArray { + fn from(columns: &[Vec]) -> Self { + let n_cols = columns.len(); + if n_cols == 0 { + return NdArray::new(&[0, 0]); + } + let n_rows = columns[0].len(); + for col in columns { + assert_eq!(col.len(), n_rows, "Column length mismatch"); + } + let shape = [n_rows, n_cols]; + let dims = NdDims::from_shape(&shape); + let stride = dims.strides()[1]; + let total = buffer_len(&shape, dims.strides()); + let mut buf = Vec64::with_capacity(total); + buf.0.resize(total, T::default()); + for (c, col) in columns.iter().enumerate() { + let start = c * stride; + buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col); + } + NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None } + } +} + +/// 2D from FloatArray columns. +impl From<&[FloatArray]> for NdArray { + fn from(columns: &[FloatArray]) -> Self { + let n_cols = columns.len(); + if n_cols == 0 { + return NdArray::new(&[0, 0]); + } + let n_rows = columns[0].data.len(); + for col in columns { + assert_eq!(col.data.len(), n_rows, "Column length mismatch"); + } + let shape = [n_rows, n_cols]; + let dims = NdDims::from_shape(&shape); + let stride = dims.strides()[1]; + let total = buffer_len(&shape, dims.strides()); + let mut buf = Vec64::with_capacity(total); + buf.0.resize(total, T::default()); + for (c, col) in columns.iter().enumerate() { + let start = c * stride; + buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col.data.as_slice()); + } + NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None } + } +} + +/// From Matrix - zero-copy, moves the Buffer straight across. The Matrix's +/// padded column stride carries through, so the resulting array reports +/// non-contiguous. Call `to_contiguous` to re-lay out compactly. +#[cfg(feature = "matrix")] +impl From for NdArray { + fn from(mat: Matrix) -> Self { + let shape = [mat.n_rows, mat.n_cols]; + let strides = [1, mat.stride]; + NdArray { + data: Arc::new(mat.data), + dims: NdDims::from_shape_and_strides(&shape, &strides), + name: mat.name, + } + } +} + +/// TryFrom Table - extracts numeric columns, converts nulls to NaN. +impl TryFrom<&Table> for NdArray { + type Error = MinarrowError; + + fn try_from(table: &Table) -> Result { + let n_cols = table.n_cols(); + let n_rows = table.n_rows; + if n_cols == 0 { + return Ok(NdArray::new(&[0, 0])); + } + + let shape = [n_rows, n_cols]; + let dims = NdDims::from_shape(&shape); + let stride = dims.strides()[1]; + let total = buffer_len(&shape, dims.strides()); + let mut buf = Vec64::with_capacity(total); + buf.0.resize(total, 0.0); + + for (col_idx, fa) in table.cols.iter().enumerate() { + let numeric = fa.array.try_num().map_err(|_| MinarrowError::TypeError { + from: "non-numeric", + to: "Float64", + message: Some(format!("column {} is not numeric", col_idx)), + })?; + let f64_arr = numeric.try_f64()?; + if f64_arr.data.len() != n_rows { + return Err(MinarrowError::ColumnLengthMismatch { + col: col_idx, + expected: n_rows, + found: f64_arr.data.len(), + }); + } + + let start = col_idx * stride; + let src = f64_arr.data.as_slice(); + let dst = &mut buf.as_mut_slice()[start..start + n_rows]; + + // Copy data, converting nulls to NaN + match f64_arr.null_mask.as_ref() { + Some(mask) => { + for i in 0..n_rows { + dst[i] = if mask.get(i) { src[i] } else { f64::NAN }; + } + } + None => dst.copy_from_slice(src), + } + } + + let name = if table.name.is_empty() { None } else { Some(table.name.clone()) }; + Ok(NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name }) + } +} + +impl TryFrom
for NdArray { + type Error = MinarrowError; + fn try_from(table: Table) -> Result { + NdArray::try_from(&table) + } +} + +// *** Debug ******************************************************* + +impl fmt::Debug for NdArray { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, "NdArray{}: {:?} [{}D, col-major]", + self.name.as_deref().map_or(String::new(), |n| format!(" '{}'", n)), + self.dims.shape(), + self.ndim(), + )?; + if self.ndim() == 0 { + write!(f, "\n{:8.4}", self.get(&[]).to_f64().unwrap_or(f64::NAN))?; + } else if self.ndim() == 2 { + let shape = self.dims.shape(); + let max_rows = shape[0].min(6); + let max_cols = shape[1].min(8); + for r in 0..max_rows { + write!(f, "\n[")?; + for c in 0..max_cols { + write!(f, " {:8.4}", self.get(&[r, c]).to_f64().unwrap_or(f64::NAN))?; + if c < max_cols - 1 { write!(f, ",")?; } + } + if shape[1] > 8 { write!(f, " ...")?; } + write!(f, " ]")?; + } + if shape[0] > 6 { write!(f, "\n...")?; } + } else if self.ndim() == 1 { + let n = self.dims.shape()[0].min(10); + write!(f, "\n[")?; + for i in 0..n { + write!(f, " {:8.4}", self.get(&[i]).to_f64().unwrap_or(f64::NAN))?; + if i < n - 1 { write!(f, ",")?; } + } + if self.dims.shape()[0] > 10 { write!(f, " ...")?; } + write!(f, " ]")?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::StringArray; + use crate::structs::bitmask::Bitmask; + + // *** Row selection and apply ************************************* + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn axis_selection_rank10() { + let shape = [3, 4, 2, 5, 3, 2, 4, 3, 2, 5]; + let len: usize = shape.iter().product(); + let data: Vec = (0..len).map(|i| i as f64).collect(); + let a = NdArray::from_slice(&data, &shape); + + // Mixed ranges, single indices, and a full range across all ten axes. + let v = a.s(nd![1..3, 0..2, 1, 2..5, .., 0..1, 1..3, 2, 0..2, 3..5]); + assert_eq!(v.shape(), &[2, 2, 3, 3, 1, 2, 2, 2]); + assert_eq!( + v.get(&[0, 0, 0, 0, 0, 0, 0, 0]), + a.get(&[1, 0, 1, 2, 0, 0, 1, 2, 0, 3]) + ); + assert_eq!( + v.get(&[1, 1, 2, 2, 0, 1, 1, 1]), + a.get(&[2, 1, 1, 4, 2, 0, 2, 2, 1, 4]) + ); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn axis_selection_runtime_rank() { + // Selections built at runtime for ranks beyond literal syntax. + let mut shape = vec![1usize; 100]; + shape[0] = 3; + shape[10] = 4; + shape[50] = 5; + shape[99] = 2; + let len: usize = shape.iter().product(); + let data: Vec = (0..len).map(|i| i as f64).collect(); + let a = NdArray::from_slice(&data, &shape); + + // Full range on every axis, then narrow three and collapse one. + let mut sels: Vec> = + shape.iter().map(|&n| Box::new(0..n) as Box).collect(); + sels[0] = Box::new(1..3); + sels[10] = Box::new(2usize); + sels[50] = Box::new(1..4); + let refs: Vec<&dyn DataSelector> = sels.iter().map(|s| s.as_ref()).collect(); + + let v = a.s(&refs); + assert_eq!(v.ndim(), 99); + assert_eq!(v.shape()[0], 2); + assert_eq!(v.shape()[49], 3); + assert_eq!(v.shape()[98], 2); + + let mut view_idx = vec![0usize; 99]; + view_idx[0] = 1; + view_idx[49] = 2; + view_idx[98] = 1; + let mut source_idx = vec![0usize; 100]; + source_idx[0] = 2; + source_idx[10] = 2; + source_idx[50] = 3; + source_idx[99] = 1; + assert_eq!(v.get(&view_idx), a.get(&source_idx)); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn axis_selection_trait() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Range keeps the axis, index collapses it. + let v = a.s(nd![1..3, 1]); + assert_eq!(v.shape(), &[2]); + assert_eq!(v.get(&[0]), 5.0); + assert_eq!(v.get(&[1]), 6.0); + // Selection composes on the view. + let sub = a.s(nd![0..3, 0..2]).s(nd![2, 0..2]); + assert_eq!(sub.shape(), &[2]); + assert_eq!(sub.get(&[0]), 3.0); + assert_eq!(sub.get(&[1]), 6.0); + assert_eq!(a.get_axis_count(), 2); + assert_eq!(sub.get_axis_count(), 1); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn row_selection_contiguous() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = a.r(1..3); + assert_eq!(v.shape(), &[2, 2]); + assert_eq!(v.get(&[0, 0]), 2.0); + assert_eq!(v.get(&[1, 1]), 6.0); + // The row alias selects a single observation. + let single = a.row(2); + assert_eq!(single.shape(), &[1, 2]); + assert_eq!(single.get(&[0, 1]), 6.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn row_selection_gathers_indices() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = a.r(&[2, 0]); + assert_eq!(v.shape(), &[2, 2]); + // Gathered rows follow selection order. + assert_eq!(v.get(&[0, 0]), 3.0); + assert_eq!(v.get(&[0, 1]), 6.0); + assert_eq!(v.get(&[1, 0]), 1.0); + assert_eq!(v.get(&[1, 1]), 4.0); + } + + #[test] + fn apply_maps_elements() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + a.set_name("m"); + let b = a.apply(|x| x * 10.0); + assert_eq!(b.get(&[1, 1]), 40.0); + assert_eq!(b.name.as_deref(), Some("m")); + // The source is untouched. + assert_eq!(a.get(&[1, 1]), 4.0); + } + + #[test] + fn apply_mut_in_place() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + a.apply_mut(|x| x + 0.5); + assert_eq!((&a).into_iter().collect::>(), vec![1.5, 2.5, 3.5]); + } + + #[cfg(feature = "matrix")] + #[test] + fn apply_mut_non_contiguous_touches_logical_only() { + // A Matrix-imported array carries stride padding. The logical walk + // mutates only real elements. + let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None); + let mut a = NdArray::from(mat); + assert!(!a.is_contiguous()); + a.apply_mut(|x| x * 2.0); + assert_eq!(a.get(&[0, 0]), 2.0); + assert_eq!(a.get(&[2, 1]), 12.0); + } + + #[cfg(feature = "views")] + #[test] + fn apply_axis_collapses_axis() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Sum each column lane (axis 0) - output shape [2]. + let col_sums = a.apply_axis(0, |lane| (&lane).into_iter().sum()); + assert_eq!(col_sums.shape(), &[2]); + assert_eq!(col_sums.get(&[0]), 6.0); + assert_eq!(col_sums.get(&[1]), 15.0); + // Sum each row lane (axis 1) - output shape [3]. + let row_sums = a.apply_axis(1, |lane| (&lane).into_iter().sum()); + assert_eq!(row_sums.shape(), &[3]); + assert_eq!(row_sums.get(&[0]), 5.0); + assert_eq!(row_sums.get(&[2]), 9.0); + } + + #[cfg(feature = "views")] + #[test] + fn apply_axis_3d() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let maxes = a.apply_axis(1, |lane| { + (&lane).into_iter().fold(f64::MIN, f64::max) + }); + assert_eq!(maxes.shape(), &[2, 4]); + // Lane over axis 1 at [0, .., 0] holds values at [0,j,0]. + assert_eq!(maxes.get(&[0, 0]), a.get(&[0, 2, 0])); + assert_eq!(maxes.get(&[1, 3]), a.get(&[1, 2, 3])); + } + + // **************************************************************** + // Construction + // **************************************************************** + + #[test] + fn new_zeroed_1d() { + let a = NdArray::::new(&[5]); + assert_eq!(a.ndim(), 1); + assert_eq!(a.shape(), &[5]); + assert_eq!(a.len(), 5); + assert!(!a.is_empty()); + for v in &a { assert_eq!(v, 0.0); } + } + + #[test] + fn new_zeroed_2d() { + let a = NdArray::::new(&[3, 4]); + assert_eq!(a.ndim(), 2); + assert_eq!(a.shape(), &[3, 4]); + assert_eq!(a.len(), 12); + for v in &a { assert_eq!(v, 0.0); } + } + + #[test] + fn new_zeroed_3d() { + let a = NdArray::::new(&[2, 3, 4]); + assert_eq!(a.ndim(), 3); + assert_eq!(a.shape(), &[2, 3, 4]); + assert_eq!(a.len(), 24); + } + + #[test] + fn new_zeroed_5d() { + let a = NdArray::::new(&[2, 3, 4, 5, 6]); + assert_eq!(a.ndim(), 5); + assert_eq!(a.len(), 720); + } + + #[test] + fn new_zeroed_6d_heap() { + let a = NdArray::::new(&[2, 3, 2, 2, 2, 2]); + assert_eq!(a.ndim(), 6); + assert_eq!(a.len(), 96); + } + + #[test] + fn new_named() { + let a = NdArray::::new_named(&[3, 3], "covariance"); + assert_eq!(a.name.as_deref(), Some("covariance")); + } + + #[test] + fn from_slice_1d() { + let data = [1.0, 2.0, 3.0, 4.0, 5.0]; + let a = NdArray::from_slice(&data, &[5]); + assert_eq!(a.len(), 5); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + } + + #[test] + fn from_slice_2d_column_major() { + let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let a = NdArray::from_slice(&data, &[3, 2]); + assert_eq!(a.get(&[0, 0]), 1.0); + assert_eq!(a.get(&[1, 0]), 2.0); + assert_eq!(a.get(&[2, 0]), 3.0); + assert_eq!(a.get(&[0, 1]), 4.0); + assert_eq!(a.get(&[1, 1]), 5.0); + assert_eq!(a.get(&[2, 1]), 6.0); + } + + #[test] + fn fill_and_ones() { + let a = NdArray::fill(&[2, 3], 7.0); + assert_eq!(a.len(), 6); + for v in &a { assert_eq!(v, 7.0); } + + let b = NdArray::::ones(&[4]); + for v in &b { assert_eq!(v, 1.0); } + } + + #[test] + fn eye_identity() { + let a = NdArray::::eye(3); + assert_eq!(a.shape(), &[3, 3]); + assert_eq!(a.get(&[0, 0]), 1.0); + assert_eq!(a.get(&[1, 1]), 1.0); + assert_eq!(a.get(&[2, 2]), 1.0); + assert_eq!(a.get(&[0, 1]), 0.0); + assert_eq!(a.get(&[1, 0]), 0.0); + } + + #[test] + fn linspace_basic() { + let a = NdArray::::linspace(0.0, 1.0, 5); + assert_eq!(a.shape(), &[5]); + assert_eq!(a.get(&[0]), 0.0); + assert_eq!(a.get(&[4]), 1.0); + assert!((a.get(&[2]) - 0.5).abs() < 1e-15); + } + + #[test] + fn arange_basic() { + let a = NdArray::arange(0.0, 0.5, 4); + assert_eq!(a.shape(), &[4]); + assert_eq!(a.get(&[0]), 0.0); + assert_eq!(a.get(&[1]), 0.5); + assert_eq!(a.get(&[2]), 1.0); + assert_eq!(a.get(&[3]), 1.5); + } + + #[test] + fn from_vec64_moves_data() { + let data: Vec64 = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into(); + let a = NdArray::from_vec64(data, &[3, 2]); + assert_eq!(a.shape(), &[3, 2]); + assert_eq!(a.get(&[2, 1]), 6.0); + } + + #[cfg(feature = "views")] + #[test] + fn iter_obs_walks_observations() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let rows: Vec<(usize, Vec)> = a + .iter_obs() + .map(|(i, v)| (i, (&v).into_iter().collect())) + .collect(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], (0, vec![1.0, 4.0])); + assert_eq!(rows[2], (2, vec![3.0, 6.0])); + } + + #[test] + fn equality_ignores_name() { + let a = NdArray::from_slice(&[1.0, 2.0], &[2]); + let mut b = a.clone(); + b.name = Some("named".to_string()); + assert_eq!(a, b); + } + + #[test] + fn name_survives_reshape_and_concat() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]); + a.name = Some("px".to_string()); + assert_eq!(a.reshape(&[2, 2]).unwrap().name.as_deref(), Some("px")); + + let mut b = NdArray::from_slice(&[5.0, 6.0], &[2]); + b.name = Some("qty".to_string()); + let joined = a.concat(b).unwrap(); + assert_eq!(joined.name.as_deref(), Some("px+qty")); + } + + // **************************************************************** + // Element access and indexing + // **************************************************************** + + #[test] + fn get_set_1d() { + let mut a = NdArray::new(&[3]); + a.set(&[0], 10.0); + a.set(&[1], 20.0); + a.set(&[2], 30.0); + assert_eq!(a.get(&[0]), 10.0); + assert_eq!(a.get(&[1]), 20.0); + assert_eq!(a.get(&[2]), 30.0); + } + + #[test] + fn tuple_index_1d() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]); + assert_eq!(a[(0,)], 10.0); + assert_eq!(a[(1,)], 20.0); + assert_eq!(a[(2,)], 30.0); + } + + #[test] + fn tuple_index_2d() { + let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let a = NdArray::from_slice(&data, &[3, 2]); + assert_eq!(a[(0, 0)], 1.0); + assert_eq!(a[(2, 0)], 3.0); + assert_eq!(a[(0, 1)], 4.0); + assert_eq!(a[(2, 1)], 6.0); + } + + #[test] + fn tuple_index_3d() { + let data: Vec = (1..=12).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 2]); + assert_eq!(a[(0, 0, 0)], 1.0); + assert_eq!(a[(1, 0, 0)], 2.0); + assert_eq!(a[(0, 1, 0)], 3.0); + assert_eq!(a[(1, 1, 0)], 4.0); + assert_eq!(a[(0, 2, 0)], 5.0); + assert_eq!(a[(1, 2, 0)], 6.0); + assert_eq!(a[(0, 0, 1)], 7.0); + assert_eq!(a[(1, 2, 1)], 12.0); + } + + #[test] + fn index_mut_2d() { + let mut a = NdArray::new(&[2, 2]); + a[(0, 0)] = 1.0; + a[(1, 0)] = 2.0; + a[(0, 1)] = 3.0; + a[(1, 1)] = 4.0; + assert_eq!(a[(0, 0)], 1.0); + assert_eq!(a[(1, 1)], 4.0); + } + + // **************************************************************** + // Iteration + // **************************************************************** + + #[test] + fn iter_1d() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals, vec![10.0, 20.0, 30.0]); + } + + #[test] + fn iter_2d_column_major_order() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + #[test] + fn iter_2d_compact_layout() { + // Compact strides place columns back to back, so iteration reads + // the buffer straight through. + let data: Vec = (1..=30).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[10, 3]); + assert_eq!(a.strides()[1], 10); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals.len(), 30); + assert_eq!(&vals[..10], &data[..10]); + assert_eq!(&vals[10..20], &data[10..20]); + assert_eq!(&vals[20..30], &data[20..30]); + } + + #[test] + fn iter_3d_column_major_order() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals.len(), 24); + assert_eq!(vals, data); + } + + #[test] + fn iter_exact_size() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let iter = (&a).into_iter(); + assert_eq!(iter.len(), 6); + } + + #[test] + fn consuming_into_iter() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let vals: Vec = a.into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0]); + } + + // **************************************************************** + // Shape introspection + // **************************************************************** + + #[test] + fn rank_zero_scalar_semantics() { + let mut a = NdArray::from_slice(&[5.0], &[]); + assert_eq!(a.ndim(), 0); + assert!(a.shape().is_empty()); + assert!(a.strides().is_empty()); + assert_eq!(a.len(), 1); + assert!(!a.is_empty()); + assert!(a.is_contiguous()); + assert_eq!(a[()], 5.0); + assert_eq!((&a).into_iter().collect::>(), vec![5.0]); + assert_eq!(Shape::shape(&a), ShapeDim::Rank0(1)); + + a[()] = 7.0; + assert_eq!(a[()], 7.0); + assert!(a.transpose().shape().is_empty()); + assert_eq!(a.reshape(&[1]).unwrap().get(&[0]), 7.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn selecting_every_axis_can_produce_a_scalar_view() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let scalar = a.slice(&[&1usize, &1usize]); + assert!(scalar.shape().is_empty()); + assert!(scalar.strides().is_empty()); + assert_eq!(scalar.len(), 1); + assert_eq!(scalar.get(&[]), 4.0); + assert_eq!((&scalar).into_iter().collect::>(), vec![4.0]); + assert_eq!(Shape::shape(&scalar), ShapeDim::Rank0(1)); + } + + #[test] + fn is_contiguous_default() { + let a = NdArray::::new(&[3, 4]); + assert!(a.is_contiguous()); + } + + #[test] + fn shape_trait_1d() { + let a = NdArray::::new(&[5]); + assert_eq!(Shape::shape(&a), ShapeDim::Rank1(5)); + } + + #[test] + fn shape_trait_2d() { + let a = NdArray::::new(&[3, 4]); + assert_eq!(Shape::shape(&a), ShapeDim::Rank2 { rows: 3, cols: 4 }); + } + + #[test] + fn shape_trait_3d() { + let a = NdArray::::new(&[2, 3, 4]); + assert_eq!(Shape::shape(&a), ShapeDim::RankN(vec![2, 3, 4])); + } + + // **************************************************************** + // NaN handling + // **************************************************************** + + #[test] + fn has_nan_false() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + assert!(!a.has_nan()); + assert_eq!(a.nan_count(), 0); + } + + #[test] + fn has_nan_true() { + let a = NdArray::from_slice(&[1.0, f64::NAN, 3.0], &[3]); + assert!(a.has_nan()); + assert_eq!(a.nan_count(), 1); + } + + #[test] + fn has_nan_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, f64::NAN, 4.0, 5.0, f64::NAN], &[3, 2]); + assert!(a.has_nan()); + assert_eq!(a.nan_count(), 2); + } + + // **************************************************************** + // 2D axis access + // **************************************************************** + + #[test] + fn col_access() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + } + + #[test] + fn col_mut_access() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + a.col_mut(0)[1] = 99.0; + assert_eq!(a.col(0), &[1.0, 99.0, 3.0]); + } + + #[test] + fn columns_access() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let cols = a.columns(); + assert_eq!(cols.len(), 2); + assert_eq!(cols[0], &[1.0, 2.0, 3.0]); + assert_eq!(cols[1], &[4.0, 5.0, 6.0]); + } + + #[test] + fn columns_mut_access() { + let mut a = NdArray::new(&[3, 2]); + { + let mut cols = a.columns_mut(); + cols[0].copy_from_slice(&[1.0, 2.0, 3.0]); + cols[1].copy_from_slice(&[4.0, 5.0, 6.0]); + } + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + } + + #[cfg(feature = "views")] + #[test] + fn obs_access() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = a.as_view(); + let obs0: Vec = (&v.obs(0)).into_iter().collect(); + let obs1: Vec = (&v.obs(1)).into_iter().collect(); + let obs2: Vec = (&v.obs(2)).into_iter().collect(); + assert_eq!(obs0, vec![1.0, 4.0]); + assert_eq!(obs1, vec![2.0, 5.0]); + assert_eq!(obs2, vec![3.0, 6.0]); + + // Single-shot .obs() on NdArray also works + assert_eq!((&a.obs(1)).into_iter().collect::>(), vec![2.0, 5.0]); + + // 3D obs returns a 2D view + let b = NdArray::from_slice(&(1..=24).map(|x| x as f64).collect::>(), &[2, 3, 4]); + let obs0_3d = b.obs(0); + assert_eq!(obs0_3d.shape(), &[3, 4]); + } + + #[test] + fn col_access_2d() { + let data: Vec = (1..=20).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[10, 2]); + assert_eq!(a.col(0).len(), 10); + assert_eq!(a.col(1).len(), 10); + assert_eq!(a.col(0), &data[..10]); + assert_eq!(a.col(1), &data[10..20]); + } + + // **************************************************************** + // BLAS compatibility + // **************************************************************** + + #[test] + fn blas_params() { + let a = NdArray::::new(&[10, 5]); + assert_eq!(a.m(), 10); + assert_eq!(a.n(), 5); + assert_eq!(a.lda(), 10); + } + + #[test] + fn blas_params_aligned_rows() { + let a = NdArray::::new(&[8, 3]); + assert_eq!(a.m(), 8); + assert_eq!(a.n(), 3); + assert_eq!(a.lda(), 8); + } + + // **************************************************************** + // Compact strides + // **************************************************************** + + #[test] + fn compact_strides_2d() { + for n_rows in 1..=20 { + let a = NdArray::::new(&[n_rows, 3]); + assert_eq!(a.strides()[1], n_rows); + assert!(a.is_contiguous()); + } + } + + #[test] + fn compact_strides_3d() { + let a = NdArray::::new(&[10, 3, 4]); + let strides = a.strides(); + assert_eq!(strides[0], 1); + assert_eq!(strides[1], 10); + assert_eq!(strides[2], 10 * 3); + assert!(a.is_contiguous()); + } + + // **************************************************************** + // Reshape and transform + // **************************************************************** + + #[test] + fn reshape_1d_to_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[6]); + let b = a.reshape(&[3, 2]).unwrap(); + assert_eq!(b.shape(), &[3, 2]); + assert_eq!(b.get(&[0, 0]), 1.0); + assert_eq!(b.get(&[1, 0]), 2.0); + assert_eq!(b.get(&[2, 0]), 3.0); + assert_eq!(b.get(&[0, 1]), 4.0); + } + + #[test] + fn reshape_size_mismatch() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + assert!(a.reshape(&[2, 2]).is_err()); + } + + #[test] + fn transpose_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let t = a.transpose(); + assert_eq!(t.shape(), &[2, 3]); + assert_eq!(t.get(&[0, 0]), 1.0); + assert_eq!(t.get(&[1, 0]), 4.0); + assert_eq!(t.get(&[0, 1]), 2.0); + assert_eq!(t.get(&[1, 1]), 5.0); + assert_eq!(t.get(&[0, 2]), 3.0); + assert_eq!(t.get(&[1, 2]), 6.0); + } + + #[test] + fn transpose_3d_reverses_axes() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let t = a.transpose(); + assert_eq!(t.shape(), &[4, 3, 2]); + assert!(t.is_contiguous()); + // Every element lands at its reversed index. + for i in 0..2 { + for j in 0..3 { + for k in 0..4 { + assert_eq!(t.get(&[k, j, i]), a.get(&[i, j, k])); + } + } + } + } + + #[test] + fn transpose_1d_copies_through() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let t = a.transpose(); + assert_eq!(t.shape(), &[3]); + assert_eq!(t, a); + } + + #[test] + fn flatten() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let flat = a.flatten(); + assert_eq!(flat.shape(), &[6]); + let vals: Vec = (&flat).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + #[test] + fn to_contiguous_noop() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let b = a.to_contiguous(); + assert_eq!(a, b); + } + + #[test] + fn fill_with_contiguous() { + let mut a = NdArray::new(&[3, 2]); + a.fill_with(42.0); + for v in &a { assert_eq!(v, 42.0); } + } + + // **************************************************************** + // From conversions + // **************************************************************** + + #[test] + fn from_f64_slice() { + let a = NdArray::from(&[1.0, 2.0, 3.0][..]); + assert_eq!(a.ndim(), 1); + assert_eq!(a.len(), 3); + assert_eq!(a[(0,)], 1.0); + } + + #[test] + fn from_vec64() { + let v: Vec64 = vec![10.0, 20.0].into_iter().collect(); + let a = NdArray::from(v); + assert_eq!(a.ndim(), 1); + assert_eq!(a[(0,)], 10.0); + assert_eq!(a[(1,)], 20.0); + } + + #[test] + fn from_column_vecs() { + let cols = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]; + let a = NdArray::from(cols.as_slice()); + assert_eq!(a.shape(), &[3, 2]); + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + } + + #[test] + fn from_float_arrays() { + let c0 = FloatArray::from_slice(&[1.0, 2.0]); + let c1 = FloatArray::from_slice(&[3.0, 4.0]); + let a = NdArray::from([c0, c1].as_slice()); + assert_eq!(a.shape(), &[2, 2]); + assert_eq!(a.col(0), &[1.0, 2.0]); + assert_eq!(a.col(1), &[3.0, 4.0]); + } + + #[test] + fn from_buffer_explicit_strides() { + let mut buf = Vec64::with_capacity(16); + buf.0.resize(16, 0.0); + buf[0] = 1.0; buf[1] = 2.0; buf[2] = 3.0; + buf[8] = 4.0; buf[9] = 5.0; buf[10] = 6.0; + let a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]); + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + } + + // **************************************************************** + // Matrix interop + // **************************************************************** + + #[cfg(feature = "matrix")] + #[test] + fn from_matrix() { + let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, Some("m".into())); + let a = NdArray::from(mat); + assert_eq!(a.shape(), &[3, 2]); + assert_eq!(a.name.as_deref(), Some("m")); + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + } + + #[cfg(feature = "matrix")] + #[test] + fn to_matrix_roundtrip() { + let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let a = NdArray::from_slice(&data, &[3, 2]); + let mat = a.to_matrix().unwrap(); + assert_eq!(mat.n_rows, 3); + assert_eq!(mat.n_cols, 2); + assert_eq!(mat.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(mat.col(1), &[4.0, 5.0, 6.0]); + } + + #[cfg(feature = "matrix")] + #[test] + fn to_matrix_non_2d_fails() { + let a = NdArray::new(&[5]); + assert!(a.to_matrix().is_err()); + } + + // **************************************************************** + // Table interop + // **************************************************************** + + fn make_numeric_table() -> Table { + let c0 = FieldArray::from_arr("x", Array::NumericArray( + NumericArray::Float64(Arc::new(FloatArray::from_slice(&[1.0, 2.0, 3.0]))) + )); + let c1 = FieldArray::from_arr("y", Array::NumericArray( + NumericArray::Float64(Arc::new(FloatArray::from_slice(&[4.0, 5.0, 6.0]))) + )); + Table::new("data".to_string(), Some(vec![c0, c1])) + } + + #[test] + fn try_from_table() { + let table = make_numeric_table(); + let a = NdArray::try_from(&table).unwrap(); + assert_eq!(a.shape(), &[3, 2]); + assert_eq!(a.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(a.col(1), &[4.0, 5.0, 6.0]); + assert_eq!(a.name.as_deref(), Some("data")); + } + + #[test] + fn try_from_table_with_nulls_converts_to_nan() { + let mut mask = Bitmask::new_set_all(3, true); + mask.set(1, false); + let arr = FloatArray::new(Buffer::from_slice(&[10.0, 0.0, 30.0]), Some(mask)); + let c0 = FieldArray::from_arr("v", Array::NumericArray( + NumericArray::Float64(Arc::new(arr)) + )); + let table = Table::new("nulls".to_string(), Some(vec![c0])); + let a = NdArray::try_from(&table).unwrap(); + assert_eq!(a.get(&[0, 0]), 10.0); + assert!(a.get(&[1, 0]).is_nan()); + assert_eq!(a.get(&[2, 0]), 30.0); + } + + #[test] + fn try_from_table_coerces_unparseable_text_to_nan() { + // TryFrom<&Table> uses the library's lenient numeric cast. Text that + // does not parse as a number coerces to nulls, which surface as NaN + // in the dense NdArray rather than failing the conversion. + let c0 = FieldArray::from_arr("name", Array::from_string32( + StringArray::from_slice(&["a", "b"]) + )); + let table = Table::new("text".to_string(), Some(vec![c0])); + let a = NdArray::try_from(&table).unwrap(); + assert_eq!(a.shape(), &[2, 1]); + assert!(a.get(&[0, 0]).is_nan()); + assert!(a.get(&[1, 0]).is_nan()); + } + + #[test] + fn to_table_roundtrip() { + let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let a = NdArray::from_slice(&data, &[3, 2]); + let fields = vec![ + Field::new("x", ArrowType::Float64, false, None), + Field::new("y", ArrowType::Float64, false, None), + ]; + let table = a.to_table(Some(fields)).unwrap(); + assert_eq!(table.n_rows(), 3); + assert_eq!(table.n_cols(), 2); + assert_eq!(table.col_names(), vec!["x", "y"]); + let col0 = table.cols[0].array.num().f64(); + assert_eq!(col0.data.as_slice(), &[1.0, 2.0, 3.0]); + } + + #[test] + fn to_table_generated_names() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let table = a.to_table(None).unwrap(); + assert_eq!(table.col_names(), vec!["col_0", "col_1"]); + } + + #[test] + fn to_table_rejects_incompatible_field_metadata() { + let wrong_dtype = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![ + Field::new("value", ArrowType::Float32, false, None), + ])); + assert!(matches!(wrong_dtype, Err(MinarrowError::TypeError { .. }))); + + let nullable = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![ + Field::new("value", ArrowType::Float64, true, None), + ])); + assert!(matches!(nullable, Err(MinarrowError::NullError { .. }))); + } + + #[test] + fn to_table_non_2d_fails() { + let a = NdArray::new(&[5]); + assert!(a.to_table(None).is_err()); + } + + #[test] + fn to_array_1d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let arr = a.to_array().unwrap(); + let f = arr.num().f64(); + assert_eq!(f.data.as_slice(), &[1.0, 2.0, 3.0]); + } + + #[test] + fn to_array_materialises_strided_logical_values() { + let a = NdArray::from_buffer( + Buffer::from_slice(&[1.0, 99.0, 2.0, 99.0, 3.0]), + &[3], + &[2], + ); + let arr = a.to_array().unwrap(); + assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]); + } + + #[test] + fn to_array_ignores_unused_backing_elements() { + let a = NdArray::from_buffer( + Buffer::from_slice(&[1.0, 2.0, 3.0, 99.0]), + &[3], + &[1], + ); + let arr = a.to_array().unwrap(); + assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]); + } + + #[test] + fn to_array_non_1d_fails() { + let a = NdArray::new(&[2, 3]); + assert!(a.to_array().is_err()); + } + + // **************************************************************** + // Concatenate + // **************************************************************** + + #[test] + fn concat_1d() { + let a = NdArray::from_slice(&[1.0, 2.0], &[2]); + let b = NdArray::from_slice(&[3.0, 4.0, 5.0], &[3]); + let c = a.concat(b).unwrap(); + assert_eq!(c.shape(), &[5]); + let vals: Vec = (&c).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + } + + #[test] + fn concat_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]); + let c = a.concat(b).unwrap(); + assert_eq!(c.shape(), &[5, 2]); + assert_eq!(c.col(0), &[1.0, 2.0, 5.0, 6.0, 7.0]); + assert_eq!(c.col(1), &[3.0, 4.0, 8.0, 9.0, 10.0]); + } + + #[test] + fn concat_dimension_mismatch_fails() { + let a = NdArray::::new(&[3, 2]); + let b = NdArray::new(&[3, 3]); + assert!(a.concat(b).is_err()); + } + + #[test] + fn concat_rank_mismatch_fails() { + let a = NdArray::::new(&[3]); + let b = NdArray::new(&[3, 2]); + assert!(a.concat(b).is_err()); + } + + // **************************************************************** + // Clone and PartialEq + // **************************************************************** + + #[test] + fn clone_and_eq() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let b = a.clone(); + assert_eq!(a, b); + } + + #[test] + fn ne_different_data() { + let a = NdArray::from_slice(&[1.0, 2.0], &[2]); + let b = NdArray::from_slice(&[1.0, 3.0], &[2]); + assert_ne!(a, b); + } + + #[test] + fn ne_different_shape() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]); + let b = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + assert_ne!(a, b); + } + + // **************************************************************** + // Debug formatting + // **************************************************************** + + #[test] + fn debug_1d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let s = format!("{:?}", a); + assert!(s.contains("[3]")); + assert!(s.contains("1D")); + } + + #[test] + fn debug_2d_named() { + let a = NdArray::::new_named(&[2, 3], "test"); + let s = format!("{:?}", a); + assert!(s.contains("'test'")); + assert!(s.contains("[2, 3]")); + assert!(s.contains("2D")); + } + + // **************************************************************** + // Edge cases + // **************************************************************** + + #[test] + fn empty_array() { + let a = NdArray::new(&[0, 5]); + assert!(a.is_empty()); + assert_eq!(a.len(), 0); + let vals: Vec = (&a).into_iter().collect(); + assert!(vals.is_empty()); + } + + #[test] + fn single_element() { + let a = NdArray::from_slice(&[42.0], &[1]); + assert_eq!(a.len(), 1); + assert_eq!(a[(0,)], 42.0); + let vals: Vec = (&a).into_iter().collect(); + assert_eq!(vals, vec![42.0]); + } + + #[test] + fn single_element_2d() { + let a = NdArray::from_slice(&[42.0], &[1, 1]); + assert_eq!(a[(0, 0)], 42.0); + } + + #[test] + fn large_array_iteration_count() { + let n = 1000; + let a = NdArray::::ones(&[n, 100]); + let count = (&a).into_iter().count(); + assert_eq!(count, n * 100); + } + + // **************************************************************** + // Bracket indexing: arr[col][row] + // **************************************************************** + + #[test] + fn bracket_index_1d() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]); + assert_eq!(a[0], [10.0]); + assert_eq!(a[2], [30.0]); + } + + #[test] + fn bracket_index_2d_column() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(a[0], [1.0, 2.0, 3.0]); + assert_eq!(a[1], [4.0, 5.0, 6.0]); + } + + #[test] + fn bracket_index_2d_chained() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // arr[col][row] + assert_eq!(a[0][0], 1.0); + assert_eq!(a[0][2], 3.0); + assert_eq!(a[1][0], 4.0); + assert_eq!(a[1][2], 6.0); + } + + #[test] + fn bracket_index_mut_2d() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + a[0][1] = 99.0; + assert_eq!(a[0][1], 99.0); + assert_eq!(a[0], [1.0, 99.0, 3.0]); + } + + // **************************************************************** + // Range indexing: arr[1..3] + // **************************************************************** + + #[test] + fn range_index_1d() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]); + assert_eq!(a[1..3], [20.0, 30.0]); + } + + #[test] + fn range_index_2d_columns() { + // Selecting a range of columns returns the exact logical data, + // since the compact layout has no padding between columns. + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let slab = &a[0..2]; + assert_eq!(slab, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let col1 = &a[1..2]; + assert_eq!(col1, &[4.0, 5.0, 6.0]); + } + + #[cfg(feature = "matrix")] + #[test] + #[should_panic(expected = "contiguous")] + fn range_index_non_contiguous_panics() { + // A Matrix-imported array carries the padded stride, so range + // indexing has no gap-free slab to return and panics. + let mat = Matrix::from_f64_unaligned( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None, + ); + let a = NdArray::from(mat); + if a.is_contiguous() { + // Padding only appears when rows are off the alignment boundary, + // which holds for 3 rows. Guard the premise. + panic!("premise failed: expected non-contiguous import"); + } + let _ = &a[0..2]; + } + + #[cfg(feature = "matrix")] + #[test] + fn to_matrix_repacks_compact_layout() { + // Compact tensor data re-lays into Matrix's padded column layout. + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let mat = a.to_matrix().unwrap(); + assert_eq!(mat.n_rows, 3); + assert_eq!(mat.n_cols, 2); + assert_eq!(mat.stride, 8); + assert_eq!(&mat.data.as_slice()[..3], &[1.0, 2.0, 3.0]); + assert_eq!(&mat.data.as_slice()[8..11], &[4.0, 5.0, 6.0]); + } + + #[cfg(feature = "matrix")] + #[test] + fn matrix_roundtrip_via_contiguous() { + // Matrix -> NdArray is zero-copy with the padded stride carried + // through. to_contiguous compacts, and to_matrix re-pads. + let mat = Matrix::from_f64_unaligned( + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None, + ); + let a = NdArray::from(mat); + assert!(!a.is_contiguous()); + assert_eq!(a.get(&[2, 1]), 6.0); + let compact = a.to_contiguous(); + assert!(compact.is_contiguous()); + assert_eq!(&compact[0..2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let back = compact.to_matrix().unwrap(); + assert_eq!(back.stride, 8); + assert_eq!(&back.data.as_slice()[8..11], &[4.0, 5.0, 6.0]); + } + + #[test] + fn range_from_index() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]); + assert_eq!(a[2..], [30.0, 40.0]); + } + + #[test] + fn range_to_index() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]); + assert_eq!(a[..2], [10.0, 20.0]); + } + + #[test] + fn range_full_index() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]); + assert_eq!(a[..], [10.0, 20.0, 30.0]); + } + + // **************************************************************** + // Slicing: arr.slice(nd![1..4, 2..5]) + // **************************************************************** + + #[cfg(feature = "views")] + #[test] + fn slice_1d_single_index() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]); + let v = a.slice(&[&1]); + assert!(v.shape().is_empty()); + assert_eq!(v[()], 20.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_1d_range() { + let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]); + let v = a.slice(&[&(1..3)]); + assert_eq!(v.shape(), &[2]); + assert_eq!(v[(0,)], 20.0); + assert_eq!(v[(1,)], 30.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_2d_row_range_single_col() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Rows 0..2 of column 1 + let v = a.slice(nd![0..2, 1]); + assert_eq!(v.shape(), &[2]); + assert_eq!(v[(0,)], 4.0); + assert_eq!(v[(1,)], 5.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_2d_single_row_col_range() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Row 1, columns 0..2 - collapses row axis + let v = a.slice(nd![1, 0..2]); + assert_eq!(v.shape(), &[2]); + // Should get row 1 values: a[(1,0)]=2.0, a[(1,1)]=5.0 + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + } + + #[cfg(feature = "views")] + #[test] + fn slice_2d_both_ranges() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Rows 0..2, columns 0..2 - sub-matrix + let v = a.slice(nd![0..2, 0..2]); + assert_eq!(v.shape(), &[2, 2]); + assert_eq!(v[(0, 0)], 1.0); + assert_eq!(v[(1, 0)], 2.0); + assert_eq!(v[(0, 1)], 4.0); + assert_eq!(v[(1, 1)], 5.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_2d_both_indices_scalar() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + // Single element as a rank-zero scalar view + let v = a.slice(nd![2, 1]); + assert!(v.shape().is_empty()); + assert_eq!(v[()], 6.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_3d_mixed() { + // 2x3x4 array + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + // All rows, column 1, slices 0..2 + let v = a.slice(nd![0..2, 1, 0..2]); + assert_eq!(v.shape(), &[2, 2]); + // a[(0,1,0)]=3, a[(1,1,0)]=4, a[(0,1,1)]=9, a[(1,1,1)]=10 + assert_eq!(v[(0, 0)], 3.0); + assert_eq!(v[(1, 0)], 4.0); + assert_eq!(v[(0, 1)], 9.0); + assert_eq!(v[(1, 1)], 10.0); + } + + #[cfg(feature = "views")] + #[test] + fn slice_with_nd_macro() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let v = a.slice(nd![0..2, 0..3, 0..4]); + assert_eq!(v.shape(), &[2, 3, 4]); + assert_eq!(v.len(), 24); + } + + #[cfg(feature = "views")] + #[test] + fn slice_preserves_data_through_iteration() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = a.slice(nd![1..3, 0..2]); + // Sub-matrix: rows 1..3, cols 0..2 + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals, vec![2.0, 3.0, 5.0, 6.0]); + } + + #[cfg(feature = "views")] + #[test] + fn slice_column_window() { + // Slice rows 2..5 from column 1 of a [10, 2] array + let data: Vec = (1..=20).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[10, 2]); + let v = a.slice(nd![2..5, 1]); + assert_eq!(v.shape(), &[3]); + assert_eq!(v[(0,)], 13.0); + assert_eq!(v[(1,)], 14.0); + assert_eq!(v[(2,)], 15.0); + } + + // **************************************************************** + // Bounds and panic contracts + // **************************************************************** + + #[test] + #[should_panic(expected = "Column index out of bounds")] + fn col_out_of_bounds_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let _ = a.col(2); + } + + #[test] + #[should_panic(expected = "indices for a 2D array")] + fn get_rank_mismatch_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let _ = a.get(&[0]); + } + + #[test] + #[should_panic(expected = "out of bounds for dim")] + fn get_index_out_of_bounds_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let _ = a.get(&[2, 0]); + } + + #[test] + #[should_panic(expected = "out of bounds for dim")] + fn set_out_of_bounds_panics() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + a.set(&[0, 5], 9.0); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + #[should_panic(expected = "expected 2 axes, got 1")] + fn slice_wrong_axis_count_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let _ = a.slice(nd![0..2]); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + #[should_panic(expected = "range 0..100 out of bounds")] + fn slice_range_out_of_bounds_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let _ = a.slice(nd![0..100, 0..2]); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + #[should_panic(expected = "index 5 out of bounds")] + fn slice_single_index_out_of_bounds_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let _ = a.slice(nd![5, 0..2]); + } + + #[test] + #[should_panic(expected = "at least 2 points")] + fn linspace_requires_two_points() { + let _ = NdArray::::linspace(0.0, 1.0, 1); + } + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn slice_one_element_index_array_collapses() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let idx: &[usize] = &[1]; + let v = a.slice(nd![idx, 0..2]); + assert_eq!(v.shape(), &[2]); + assert_eq!(v[(0,)], 2.0); + assert_eq!(v[(1,)], 5.0); + } + + #[test] + fn get_unchecked_matches_get() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + for i in 0..2 { + for j in 0..3 { + for k in 0..4 { + let idx = [i, j, k]; + // SAFETY: indices are within shape + assert_eq!(unsafe { a.get_unchecked(&idx) }, a.get(&idx)); + } + } + } + } + + // **************************************************************** + // Concatenate correctness + // **************************************************************** + + #[test] + fn concat_3d_interleaves_axis0() { + let da: Vec = (1..=8).map(|x| x as f64).collect(); + let db: Vec = (9..=16).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&da, &[2, 2, 2]); + let b = NdArray::from_slice(&db, &[2, 2, 2]); + let c = a.clone().concat(b.clone()).unwrap(); + assert_eq!(c.shape(), &[4, 2, 2]); + for i in 0..4 { + for j in 0..2 { + for k in 0..2 { + let expected = if i < 2 { + a.get(&[i, j, k]) + } else { + b.get(&[i - 2, j, k]) + }; + assert_eq!(c.get(&[i, j, k]), expected); + } + } + } + } + + #[test] + fn concat_non_contiguous_operand() { + // Row-major strides on the first operand exercise the general + // interleave path. + let a = NdArray::from_buffer( + Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]), + &[2, 2], + &[2, 1], + ); + assert!(!a.is_contiguous()); + let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]); + let c = a.concat(b).unwrap(); + assert_eq!(c.shape(), &[4, 2]); + assert_eq!(c.col(0), &[1.0, 3.0, 5.0, 6.0]); + assert_eq!(c.col(1), &[2.0, 4.0, 7.0, 8.0]); + } + + // **************************************************************** + // Copy-on-write and strided mutation + // **************************************************************** + + #[cfg(feature = "views")] + #[test] + fn set_after_view_copy_on_write() { + let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let v = a.as_view(); + a.set(&[0, 0], 99.0); + // The write detached the array's buffer, so the view keeps the + // original value. + assert_eq!(a.get(&[0, 0]), 99.0); + assert_eq!(v.get(&[0, 0]), 1.0); + } + + #[test] + fn fill_with_non_contiguous_logical_only() { + // Padded column stride leaves gaps between columns in the buffer. + let mut buf = Vec64::with_capacity(11); + buf.0.resize(11, 0.0); + let mut a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]); + assert!(!a.is_contiguous()); + a.fill_with(7.0); + for v in &a { assert_eq!(v, 7.0); } + // Padding between the columns stays untouched. + assert_eq!(a.as_slice()[3], 0.0); + assert_eq!(a.as_slice()[7], 0.0); + } + + // **************************************************************** + // Range shorthand on 2D + // **************************************************************** + + #[test] + fn range_from_index_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(&a[1..], &[4.0, 5.0, 6.0]); + } + + #[test] + fn range_to_index_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(&a[..1], &[1.0, 2.0, 3.0]); + } + + #[test] + fn range_full_index_2d() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(&a[..], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + // **************************************************************** + // Row window and size estimate + // **************************************************************** + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn row_selection_single_index_window() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert_eq!(a.n_obs(), 3); + let v = a.r(1usize); + assert_eq!(v.shape(), &[1, 2]); + assert_eq!(v.get(&[0, 0]), 2.0); + assert_eq!(v.get(&[0, 1]), 5.0); + // A window keeps the full source buffer rather than gathering a copy. + assert_eq!(v.source.len(), a.len()); + // SAFETY: indices are within the window shape + assert_eq!(unsafe { v.get_unchecked(&[0, 1]) }, 5.0); + } + + #[cfg(feature = "size")] + #[test] + fn est_bytes_covers_buffer() { + use crate::traits::byte_size::ByteSize; + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + assert!(a.est_bytes() >= 6 * size_of::()); + } +} diff --git a/src/structs/views/chunked/super_ndarray_view.rs b/src/structs/views/chunked/super_ndarray_view.rs new file mode 100644 index 0000000..610b504 --- /dev/null +++ b/src/structs/views/chunked/super_ndarray_view.rs @@ -0,0 +1,655 @@ +// Copyright 2025 Peter Garfield Bower +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! # **SuperNdArrayV** - *Window over N-dimensional batches* +//! +//! `SuperNdArrayV` is a borrowed view over an arbitrary +//! `[offset .. offset + len)` axis-0 window of a [`SuperNdArray`]. The window +//! may span multiple underlying batches while presenting one continuous +//! logical range without copying their data. +//! +//! ## Role +//! - Keeps a zero-copy window over independently landed batches, such as a +//! time range spanning several sensor or telemetry arrivals. +//! - The N-dimensional counterpart of [`SuperArrayV`](crate::SuperArrayV). +//! +//! ## Interop +//! - Constructed by [`SuperNdArray::slice`] or [`From`]. +//! - Materialises to a contiguous [`NdArray`] via [`Consolidate`]. +//! +//! ## Invariants +//! - `slices` are ordered, non-overlapping axis-0 windows sharing rank and +//! trailing shape. +//! - `n_obs` is the logical axis-0 observation count of this view. + +use std::fmt; + +use crate::enums::error::MinarrowError; +use crate::enums::shape_dim::ShapeDim; +use crate::structs::chunked::super_ndarray::SuperNdArray; +use crate::structs::ndarray::NdArray; +#[cfg(feature = "select")] +use crate::structs::ndarray::gather_obs_impl; +use crate::structs::views::ndarray_view::NdArrayV; +use crate::traits::concatenate::Concatenate; +#[cfg(feature = "select")] +use crate::traits::selection::{AxisSelection, DataSelector, RowSelection}; +use crate::traits::consolidate::Consolidate; +use crate::traits::shape::Shape; +use crate::traits::type_unions::Float; +use crate::Vec64; + +/// Borrowed view over an arbitrary `[offset .. offset + len)` axis-0 window +/// of a [`SuperNdArray`], spanning batch boundaries without copying. +/// +/// ## Fields +/// - `slices`: constituent [`NdArrayV`] windows spanning the range, ordered +/// and sharing rank and trailing shape. +/// - Rank, trailing shape, and the parent's name are cached so an empty +/// window keeps its dimensionality and identity. +#[derive(Clone)] +pub struct SuperNdArrayV { + pub slices: Vec>, + ndim: usize, + inner_shape: Vec, + name: String, +} + +impl SuperNdArrayV { + /// Assemble from ordered axis-0 window slices. Panics if the slices + /// disagree on rank or trailing shape. + pub fn from_slices( + slices: Vec>, + ndim: usize, + inner_shape: Vec, + name: String, + ) -> Self { + for (i, s) in slices.iter().enumerate() { + assert_eq!( + s.ndim(), ndim, + "SuperNdArrayV: slice {} has rank {} but expected {}", i, s.ndim(), ndim + ); + assert_eq!( + &s.shape()[1..], inner_shape.as_slice(), + "SuperNdArrayV: slice {} inner shape mismatch", i + ); + } + SuperNdArrayV { slices, ndim, inner_shape, name } + } + + /// The parent array's name. + #[inline] + pub fn name(&self) -> &str { &self.name } + + /// Number of constituent slices. + #[inline] + pub fn n_slices(&self) -> usize { self.slices.len() } + + /// Shared rank. + #[inline] + pub fn ndim(&self) -> usize { self.ndim } + + /// Dimensions shared across all slices i.e. shape[1..]. + #[inline] + pub fn inner_shape(&self) -> &[usize] { &self.inner_shape } + + /// Total axis-0 observations across all slices. + #[inline] + pub fn n_obs(&self) -> usize { + self.slices.iter().map(|s| s.shape()[0]).sum() + } + + /// Total logical elements across all slices. + #[inline] + pub fn len(&self) -> usize { + self.slices.iter().map(|s| s.len()).sum() + } + + #[inline] + pub fn is_empty(&self) -> bool { self.len() == 0 } + + /// Logical shape as if consolidated. Axis 0 is the sum across slices. + pub fn shape(&self) -> Vec { + let mut s = vec![self.n_obs()]; + s.extend_from_slice(&self.inner_shape); + s + } + + /// Iterator over the constituent slice views. + #[inline] + pub fn chunks(&self) -> impl Iterator> { + self.slices.iter() + } + + /// Iterate values in the column-major order of the logical consolidated + /// view without materialising it. Unlike `IntoIterator`, this interleaves + /// corresponding axis-0 runs across slices. + pub fn iter_logical(&self) -> impl Iterator + '_ { + let n_runs: usize = self.inner_shape.iter().product(); + (0..n_runs).flat_map(move |run| { + self.slices.iter().flat_map(move |slice| slice.iter_axis0_run(run)) + }) + } + + /// Returns a sub-window of this view over `[offset .. offset + len)` + /// axis-0 observations. Zero-copy - the new view narrows each + /// constituent slice as needed. + pub fn slice(&self, mut offset: usize, mut len: usize) -> Self { + assert!( + offset + len <= self.n_obs(), + "SuperNdArrayV::slice: window [{}, {}) out of bounds (n_obs {})", + offset, offset + len, self.n_obs() + ); + + let mut slices = Vec::new(); + for view in &self.slices { + let base_obs = view.shape()[0]; + if offset >= base_obs { + offset -= base_obs; + continue; + } + + let take = (base_obs - offset).min(len); + let mut window_shape = vec![take]; + window_shape.extend_from_slice(&view.shape()[1..]); + slices.push(NdArrayV::new( + view.source.clone(), + view.offset + offset * view.strides()[0], + &window_shape, + view.strides(), + )); + + len -= take; + if len == 0 { + break; + } + offset = 0; + } + + SuperNdArrayV { + slices, + ndim: self.ndim, + inner_shape: self.inner_shape.clone(), + name: self.name.clone(), + } + } + + /// Zero-copy view of a single observation (axis-0 element), resolving + /// which slice contains it. Returns an (N-1)-dimensional view. + pub fn obs(&self, mut idx: usize) -> NdArrayV { + for slice in &self.slices { + let n = slice.shape()[0]; + if idx < n { + return slice.obs(idx); + } + idx -= n; + } + panic!("SuperNdArrayV::obs: index out of bounds (n_obs {})", self.n_obs()); + } + + /// Get element by global N-dimensional index. The first index is the + /// global axis-0 position across slices. + pub fn get(&self, indices: &[usize]) -> T { + let mut local = indices.to_vec(); + for slice in &self.slices { + let n = slice.shape()[0]; + if local[0] < n { + return slice.get(&local); + } + local[0] -= n; + } + panic!("SuperNdArrayV::get: index out of bounds (n_obs {})", self.n_obs()); + } + + /// Apply a function to every logical element, materialising a new + /// compact [`NdArray`] with this window's shape. + pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray { + self.clone().consolidate().apply(f) + } +} + +// *** Axis selection: view.s(nd![1..4, 2]) ************************ + +/// Selection across every axis at once over a SuperNdArrayV window. The +/// axis-0 selection narrows the window, and trailing-axis selections +/// narrow each slice. Zero-copy. An axis-0 single index keeps the +/// dimension as a one-observation window - use `obs` to collapse. +/// The resulting rank and trailing shape derive from the selections, +/// so an empty window keeps its dimensionality. +#[cfg(feature = "select")] +impl AxisSelection for SuperNdArrayV { + type View = SuperNdArrayV; + + fn s(&self, selection: &[&dyn DataSelector]) -> SuperNdArrayV { + assert_eq!( + selection.len(), self.ndim, + "s(): expected {} axes, got {}", self.ndim, selection.len() + ); + let (start, end, _) = selection[0].resolve_axis(self.n_obs()); + let window = self.slice(start, end - start); + if self.ndim == 1 { + return window; + } + + let inner = &selection[1..]; + let mut inner_shape = Vec::new(); + for (d, sel) in inner.iter().enumerate() { + let (start, end, collapse) = sel.resolve_axis(window.inner_shape()[d]); + if !collapse { + inner_shape.push(end - start); + } + } + let ndim = 1 + inner_shape.len(); + + let slices: Vec> = window + .slices + .iter() + .map(|sv| { + let full = 0..sv.shape()[0]; + let mut refs: Vec<&dyn DataSelector> = vec![&full]; + refs.extend_from_slice(inner); + sv.slice(&refs) + }) + .collect(); + SuperNdArrayV::from_slices(slices, ndim, inner_shape, self.name.clone()) + } + + fn get_axis_count(&self) -> usize { + self.ndim() + } +} + +// *** Row selection: view.r(0..10) ******************************** + +/// Axis-0 observation selection over a SuperNdArrayV window. Contiguous ranges +/// narrow the window zero-copy. Index arrays gather into one owned batch +/// wrapped in a single-slice view. +#[cfg(feature = "select")] +impl RowSelection for SuperNdArrayV { + type View = SuperNdArrayV; + + fn r(&self, selection: S) -> SuperNdArrayV { + if self.slices.is_empty() { + return SuperNdArrayV::from_slices( + Vec::new(), + self.ndim, + self.inner_shape.clone(), + self.name.clone(), + ); + } + let indices = selection.resolve_indices(self.n_obs()); + if selection.is_contiguous() { + let start = indices.first().copied().unwrap_or(0); + return self.slice(start, indices.len()); + } + let gathered = gather_obs_impl( + &indices, + &self.shape(), + Some(self.name.clone()), + |idx| self.get(idx), + ); + SuperNdArrayV::from_slices( + vec![NdArrayV::from_ndarray(gathered)], + self.ndim, + self.inner_shape.clone(), + self.name.clone(), + ) + } + + fn get_row_count(&self) -> usize { + self.n_obs() + } +} + +// *** IntoIterator ************************************************ + +/// Iterating a SuperNdArrayV walks each slice in sequence, with column-major +/// order inside each slice. Use [`SuperNdArrayV::iter_logical`] for the +/// column-major order of the consolidated logical view. +impl<'a, T: Float> IntoIterator for &'a SuperNdArrayV { + type Item = T; + type IntoIter = Box + 'a>; + + fn into_iter(self) -> Self::IntoIter { + Box::new(self.slices.iter().flat_map(|s| s.into_iter())) + } +} + +// *** Trait implementations *************************************** + +impl Shape for SuperNdArrayV { + fn shape(&self) -> ShapeDim { + let obs = self.n_obs(); + match self.ndim { + 0 | 1 => ShapeDim::Rank1(obs), + 2 => ShapeDim::Rank2 { rows: obs, cols: self.inner_shape[0] }, + _ => { + let mut full = vec![obs]; + full.extend_from_slice(&self.inner_shape); + ShapeDim::RankN(full) + } + } + } +} + +impl Consolidate for SuperNdArrayV { + type Output = NdArray; + + /// Materialise the window into a contiguous compact [`NdArray`], + /// interleaving each slice's axis-0 rows one column at a time. + /// An empty window keeps its rank, trailing shape, and name. + fn consolidate(self) -> NdArray { + if self.slices.is_empty() { + let mut result = if self.ndim == 0 { + NdArray::new(&[0]) + } else { + NdArray::from_slice(&[], &self.shape()) + }; + result.name = Some(self.name); + return result; + } + + let full_shape = self.shape(); + let total_obs = full_shape[0]; + let n_cols: usize = self.inner_shape.iter().product::(); + + // Each slice's iterator yields column-major logical values, so its + // column runs arrive in order and interleave by column index. + let per_slice: Vec<(usize, Vec)> = self + .slices + .iter() + .map(|s| (s.shape()[0], s.into_iter().collect())) + .collect(); + + let mut flat: Vec64 = Vec64::with_capacity(total_obs * n_cols); + for c in 0..n_cols { + for (obs, elems) in &per_slice { + flat.extend_from_slice(&elems[c * obs..(c + 1) * obs]); + } + } + let mut result = NdArray::from_slice(&flat, &full_shape); + result.name = Some(self.name); + result + } +} + +impl Concatenate for SuperNdArrayV { + /// Concatenates two SuperNdArrayV windows along axis 0 by appending the other + /// view's slices. Zero-copy - both views' slices carry across. + fn concat(mut self, other: Self) -> Result { + if self.slices.is_empty() { + return Ok(other); + } + if other.slices.is_empty() { + return Ok(self); + } + if self.ndim != other.ndim || self.inner_shape != other.inner_shape { + return Err(MinarrowError::IncompatibleTypeError { + from: "SuperNdArrayV", + to: "SuperNdArrayV", + message: Some(format!( + "shape {:?} vs {:?}", self.shape(), other.shape() + )), + }); + } + self.slices.extend(other.slices); + Ok(self) + } +} + +/// Logical equality over shape and values in logical order. Slice +/// boundaries do not affect equality. +impl PartialEq for SuperNdArrayV { + fn eq(&self, other: &Self) -> bool { + if self.ndim != other.ndim + || self.inner_shape != other.inner_shape + || self.n_obs() != other.n_obs() + { + return false; + } + let a = self.clone().consolidate(); + let b = other.clone().consolidate(); + (&a).into_iter().eq((&b).into_iter()) + } +} + +impl fmt::Debug for SuperNdArrayV { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "SuperNdArrayV: {} slices, {}D, shape {:?}, {} elements", + self.n_slices(), + self.ndim, + self.shape(), + self.len() + ) + } +} + +/// SuperNdArray -> SuperNdArrayV conversion. Each batch becomes a full +/// axis-0 slice view, keeping the parent batches alive through each +/// batch's shared internal buffer. +impl From> for SuperNdArrayV { + fn from(super_nd: SuperNdArray) -> Self { + let ndim = super_nd.ndim(); + let inner_shape = super_nd.inner_shape().to_vec(); + let slices: Vec> = super_nd + .batches + .iter() + .map(|b| NdArrayV::from_ndarray(b.clone())) + .collect(); + SuperNdArrayV { slices, ndim, inner_shape, name: super_nd.name } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn two_batch_2d() -> SuperNdArray { + SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0, 10.0, 20.0], &[2, 2]), + NdArray::from_slice(&[3.0, 4.0, 5.0, 30.0, 40.0, 50.0], &[3, 2]), + ], + "data", + ) + } + + #[test] + fn full_view_from_super() { + let snd = two_batch_2d(); + let v = SuperNdArrayV::from(snd); + assert_eq!(v.n_slices(), 2); + assert_eq!(v.n_obs(), 5); + assert_eq!(v.shape(), vec![5, 2]); + assert_eq!(v.get(&[0, 0]), 1.0); + assert_eq!(v.get(&[2, 0]), 3.0); + assert_eq!(v.get(&[4, 1]), 50.0); + } + + #[test] + fn window_spans_batches() { + let snd = two_batch_2d(); + // Rows 1..4 span the batch boundary at row 2. + let v = snd.slice(1, 3); + assert_eq!(v.n_slices(), 2); + assert_eq!(v.n_obs(), 3); + assert_eq!(v.get(&[0, 0]), 2.0); + assert_eq!(v.get(&[1, 0]), 3.0); + assert_eq!(v.get(&[2, 1]), 40.0); + } + + #[test] + fn sub_window() { + let snd = two_batch_2d(); + let v = snd.slice(0, 5); + let sub = v.slice(1, 3); + assert_eq!(sub.n_obs(), 3); + assert_eq!(sub.get(&[0, 0]), 2.0); + assert_eq!(sub.get(&[2, 0]), 4.0); + } + + #[test] + fn obs_across_boundary() { + let snd = two_batch_2d(); + let v = snd.slice(0, 5); + let o = v.obs(3); + assert_eq!(o.shape(), &[2]); + assert_eq!(o.get(&[0]), 4.0); + assert_eq!(o.get(&[1]), 40.0); + } + + #[test] + fn consolidate_window() { + let snd = two_batch_2d(); + let v = snd.slice(1, 3); + let nd = v.consolidate(); + assert_eq!(nd.shape(), &[3, 2]); + assert!(nd.is_contiguous()); + assert_eq!(nd.col(0), &[2.0, 3.0, 4.0]); + assert_eq!(nd.col(1), &[20.0, 30.0, 40.0]); + } + + #[test] + fn iteration_crosses_slices() { + let snd = SuperNdArray::from_batches( + vec![ + NdArray::from_slice(&[1.0, 2.0], &[2]), + NdArray::from_slice(&[3.0], &[1]), + ], + "1d", + ); + let v = snd.slice(0, 3); + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0]); + } + + #[test] + fn logical_iteration_interleaves_slices_by_axis_zero_run() { + let snd = two_batch_2d(); + let view = snd.slice(1, 3); + let batch_first: Vec = (&view).into_iter().collect(); + assert_eq!(batch_first, vec![2.0, 20.0, 3.0, 4.0, 30.0, 40.0]); + + let logical: Vec = view.iter_logical().collect(); + assert_eq!(logical, vec![2.0, 3.0, 4.0, 20.0, 30.0, 40.0]); + assert_eq!( + logical, + view.clone().consolidate().into_iter().collect::>() + ); + } + + #[test] + fn eq_ignores_slice_boundaries() { + let snd = two_batch_2d(); + let whole = snd.slice(0, 5); + let single = SuperNdArrayV::from(SuperNdArray::from_batches( + vec![NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0, 30.0, 40.0, 50.0], + &[5, 2], + )], + "one", + )); + assert_eq!(whole, single); + } + + #[cfg(feature = "select")] + #[test] + fn row_selection_on_window() { + let snd = two_batch_2d(); + let v = snd.slice(0, 5); + // Contiguous sub-selection narrows zero-copy. + let sub = v.r(1..4); + assert_eq!(sub.n_obs(), 3); + assert_eq!(sub.get(&[0, 0]), 2.0); + // Index selection gathers in order. + let picked = v.r(&[4, 0]); + assert_eq!(picked.n_slices(), 1); + assert_eq!(picked.get(&[0, 1]), 50.0); + assert_eq!(picked.get(&[1, 0]), 1.0); + } + + #[test] + fn apply_materialises_window() { + let snd = two_batch_2d(); + let out = snd.slice(1, 3).apply(|x| x * 2.0); + assert_eq!(out.shape(), &[3, 2]); + assert_eq!(out.get(&[0, 0]), 4.0); + assert_eq!(out.get(&[2, 1]), 80.0); + } + + #[test] + fn concat_appends_slices() { + let snd = two_batch_2d(); + let a = snd.slice(0, 2); + let b = snd.slice(2, 3); + let joined = a.concat(b).unwrap(); + assert_eq!(joined.n_obs(), 5); + assert_eq!(joined.get(&[4, 1]), 50.0); + } + + #[test] + fn window_3d_spans_boundary() { + // Batch A holds column-major values 1..=8, batch B holds 9..=16. + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 2, 2]); + let b = NdArray::from_slice( + &[9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0], + &[2, 2, 2], + ); + let snd = SuperNdArray::from_batches(vec![a, b], "cube"); + // Rows 1..3 span the batch boundary at row 2. + let v = snd.slice(1, 2); + assert_eq!(v.n_slices(), 2); + assert_eq!(v.shape(), vec![2, 2, 2]); + assert_eq!(v.get(&[0, 0, 0]), 2.0); + assert_eq!(v.get(&[0, 1, 0]), 4.0); + assert_eq!(v.get(&[0, 0, 1]), 6.0); + assert_eq!(v.get(&[0, 1, 1]), 8.0); + assert_eq!(v.get(&[1, 0, 0]), 9.0); + assert_eq!(v.get(&[1, 1, 0]), 11.0); + assert_eq!(v.get(&[1, 0, 1]), 13.0); + assert_eq!(v.get(&[1, 1, 1]), 15.0); + + let nd = v.consolidate(); + assert_eq!(nd.shape(), &[2, 2, 2]); + assert!(nd.is_contiguous()); + assert_eq!(nd.get(&[0, 0, 0]), 2.0); + assert_eq!(nd.get(&[1, 0, 0]), 9.0); + assert_eq!(nd.get(&[0, 1, 1]), 8.0); + assert_eq!(nd.get(&[1, 1, 1]), 15.0); + } + + #[test] + fn consolidate_zero_trailing_dim() { + let snd = SuperNdArray::from_batches( + vec![NdArray::::from_slice(&[], &[2, 0])], + "hollow", + ); + let v = snd.slice(0, 2); + assert_eq!(v.n_obs(), 2); + let nd = v.consolidate(); + assert_eq!(nd.shape(), &[2, 0]); + assert_eq!(nd.len(), 0); + } + + #[test] + fn empty_view_consolidate_keeps_identity() { + let v = SuperNdArrayV::::from_slices(Vec::new(), 2, vec![3], "empty".to_string()); + let nd = v.consolidate(); + assert_eq!(nd.ndim(), 2); + assert_eq!(nd.shape(), &[0, 3]); + assert_eq!(nd.name.as_deref(), Some("empty")); + } +} diff --git a/src/structs/views/ndarray_view.rs b/src/structs/views/ndarray_view.rs new file mode 100644 index 0000000..3e58435 --- /dev/null +++ b/src/structs/views/ndarray_view.rs @@ -0,0 +1,932 @@ +//! # **NdArrayV** - *Zero-copy view into an NdArray* +//! +//! Holds a clone of the parent `NdArray`, kept alive through the array's +//! shared internal buffer, plus its own offset and dimension metadata. +//! Views can have different shapes and strides from the parent, enabling +//! slicing, axis selection, transposition, and axis permutation without +//! copying data. + +use std::fmt; +use std::ops::Index; + +use crate::enums::error::MinarrowError; +use crate::enums::shape_dim::ShapeDim; +use crate::structs::ndarray::{NdArray, NdArrayIter, NdDims, offset_of_impl}; +#[cfg(feature = "select")] +use crate::structs::ndarray::gather_obs_impl; +#[cfg(feature = "select")] +use crate::traits::selection::{AxisSelection, DataSelector, RowSelection}; +#[cfg(feature = "select")] +use std::ops::Range; +#[cfg(feature = "dlpack")] +use crate::ffi::dlpack::{ + export_view_to_dlpack, export_view_to_dlpack_versioned, DLPackTensor, DLPackTensorVersioned, +}; +use crate::traits::shape::Shape; +use crate::traits::type_unions::Float; +use crate::Vec64; + +#[cfg(feature = "matrix")] +use crate::structs::matrix::Matrix; + +/// Zero-copy view into an [`NdArray`]. +/// +/// Holds a clone of the parent `NdArray`, kept alive through the array's +/// shared internal buffer, with its own offset and dimension metadata. +/// This enables slicing, axis selection, transposition, and axis +/// permutation without copying the underlying data. +#[derive(Clone)] +pub struct NdArrayV { + pub(crate) source: NdArray, + pub(crate) offset: usize, + pub(crate) dims: NdDims, +} + +impl NdArrayV { + /// Create a view over an NdArray with the given offset and dimensions. + /// Panics when the rank metadata or reachable buffer span is invalid. + pub fn new(source: NdArray, offset: usize, shape: &[usize], strides: &[usize]) -> Self { + Self::try_new(source, offset, shape, strides) + .unwrap_or_else(|e| panic!("NdArrayV::new: {}", e)) + } + + /// Checked view construction over an NdArray. + /// + /// Validates rank metadata, arithmetic overflow, and that every logical + /// element reachable through `shape` and `strides` lies in the source + /// buffer. A shape with zero axes (`shape == &[]`) denotes one scalar + /// value. A shape containing a zero-length axis, such as `&[0]`, denotes + /// no values and may use an offset one past the buffer end. + pub fn try_new( + source: NdArray, + offset: usize, + shape: &[usize], + strides: &[usize], + ) -> Result { + if shape.len() != strides.len() { + return Err(MinarrowError::ShapeError { + message: format!( + "view shape rank {} does not match strides rank {}", + shape.len(), strides.len() + ), + }); + } + + let span = if shape.contains(&0) { + 0 + } else { + let max_offset = shape.iter().zip(strides.iter()).try_fold( + 0usize, + |acc, (&dim, &stride)| { + (dim - 1) + .checked_mul(stride) + .and_then(|extent| acc.checked_add(extent)) + }, + ).ok_or_else(|| MinarrowError::ShapeError { + message: format!("view shape {:?} with strides {:?} overflows", shape, strides), + })?; + max_offset.checked_add(1).ok_or_else(|| MinarrowError::ShapeError { + message: format!("view shape {:?} with strides {:?} overflows", shape, strides), + })? + }; + let end = offset.checked_add(span).ok_or_else(|| MinarrowError::ShapeError { + message: format!("view offset {} plus span {} overflows", offset, span), + })?; + let buffer_len = source.data.as_slice().len(); + if end > buffer_len || (span == 0 && offset > buffer_len) { + return Err(MinarrowError::IndexError(format!( + "view span [{}, {}) exceeds source buffer length {}", + offset, end, buffer_len + ))); + } + + Ok(NdArrayV { + source, + offset, + dims: NdDims::from_shape_and_strides(shape, strides), + }) + } + + /// Create a full view over an NdArray with the same shape and strides. + pub fn from_ndarray(source: NdArray) -> Self { + let dims = source.dims.clone(); + NdArrayV { source, offset: 0, dims } + } + + /// Number of dimensions. + #[inline] + pub fn ndim(&self) -> usize { self.dims.ndim() } + + /// Shape as a slice. + #[inline] + pub fn shape(&self) -> &[usize] { self.dims.shape() } + + /// Strides as a slice. + #[inline] + pub fn strides(&self) -> &[usize] { self.dims.strides() } + + /// Total logical element count. + #[inline] + pub fn len(&self) -> usize { self.dims.len() } + + /// True if empty. + #[inline] + pub fn is_empty(&self) -> bool { self.len() == 0 } + + /// Get element by N-dimensional index. + #[inline] + pub fn get(&self, indices: &[usize]) -> T { + let off = self.offset_of(indices); + self.source.data.as_slice()[off] + } + + /// Like `get`, but skips bounds checks. + /// + /// # Safety + /// The caller guarantees each index is within its dimension. + #[inline(always)] + pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T { + let strides = self.dims.strides(); + let mut off = self.offset; + for d in 0..indices.len() { + off += indices[d] * strides[d]; + } + // SAFETY: in-bounds indices produce an in-bounds flat offset. + unsafe { *self.source.data.as_slice().get_unchecked(off) } + } + + /// Compute flat offset. + #[inline] + fn offset_of(&self, indices: &[usize]) -> usize { + self.offset + offset_of_impl(indices, self.dims.shape(), self.dims.strides()) + } + + /// Immutable column slice for 2D views. + /// Panics if ndim != 2, the column index is out of bounds, or axis 0 + /// is not unit-stride i.e. column elements are not contiguous, as after + /// `transpose`, `permute_axes`, or `swap_axes`. Materialise with + /// `to_ndarray()` first for those. + #[inline] + pub fn col(&self, col: usize) -> &[T] { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "col() requires a 2D view"); + assert!(col < shape[1], "Column index out of bounds"); + assert_eq!( + self.dims.strides()[0], 1, + "col() requires unit stride on axis 0; materialise with to_ndarray() first" + ); + let stride = self.dims.strides()[1]; + let start = self.offset + col * stride; + &self.source.data.as_slice()[start..start + shape[0]] + } + + /// All columns as slices. 2D only. + /// Panics if ndim != 2 or axis 0 is not unit-stride. + pub fn columns(&self) -> Vec<&[T]> { + let shape = self.dims.shape(); + assert_eq!(shape.len(), 2, "columns() requires a 2D view"); + assert_eq!( + self.dims.strides()[0], 1, + "columns() requires unit stride on axis 0; materialise with to_ndarray() first" + ); + let stride = self.dims.strides()[1]; + let n_rows = shape[0]; + let buf = self.source.data.as_slice(); + (0..shape[1]) + .map(|c| &buf[self.offset + c * stride..self.offset + c * stride + n_rows]) + .collect() + } + + // *** BLAS/LAPACK compatibility (2D) ************************** + + /// BLAS row count. 2D only. + #[inline] + pub fn m(&self) -> i32 { + assert_eq!(self.ndim(), 2, "m() requires a 2D view"); + self.dims.shape()[0] as i32 + } + + /// BLAS column count. 2D only. + #[inline] + pub fn n(&self) -> i32 { + assert_eq!(self.ndim(), 2, "n() requires a 2D view"); + self.dims.shape()[1] as i32 + } + + /// BLAS leading dimension. 2D only. + /// Panics unless axis 0 is unit-stride - BLAS requires column + /// elements to be contiguous, which a transposed or permuted view + /// does not satisfy. Materialise with `to_ndarray()` first. + #[inline] + pub fn lda(&self) -> i32 { + assert_eq!(self.ndim(), 2, "lda() requires a 2D view"); + assert_eq!( + self.dims.strides()[0], 1, + "lda() requires unit stride on axis 0; materialise with to_ndarray() first" + ); + self.dims.strides()[1] as i32 + } + + // *** Slicing ************************************************* + + /// Zero-copy view of a single observation (axis-0 element). + /// + /// Returns an (N-1)-dimensional view. For a 2D view with shape + /// `[n, m]`, returns a 1D view of shape `[m]`. For 3D `[n, m, k]`, + /// returns 2D `[m, k]`. Requires rank 2 or higher - a 1D view has no + /// sub-array observations, so scalar access goes through `get(&[i])`. + pub fn obs(&self, idx: usize) -> NdArrayV { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + assert!( + shape.len() >= 2, + "obs() requires a 2D or higher view, use get(&[i]) for scalar access on 1D" + ); + assert!(idx < shape[0], "obs: index {} out of bounds for axis 0 (size {})", idx, shape[0]); + + let new_offset = self.offset + idx * strides[0]; + NdArrayV::new(self.source.clone(), new_offset, &shape[1..], &strides[1..]) + } + + /// Slice this view, producing a sub-view. Each axis takes any + /// [`DataSelector`] - a single index collapses that dimension, and a + /// contiguous range keeps it. Zero-copy - shares the same backing + /// buffer, just adjusts offset and dims. + #[cfg(feature = "select")] + pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + assert_eq!( + selection.len(), shape.len(), + "slice(): expected {} axes, got {}", shape.len(), selection.len() + ); + + let mut new_offset = self.offset; + let mut new_shape = Vec::with_capacity(shape.len()); + let mut new_strides = Vec::with_capacity(shape.len()); + + for (d, sel) in selection.iter().enumerate() { + let (start, end, collapse) = sel.resolve_axis(shape[d]); + assert!( + end <= shape[d], + "slice(): end {} out of bounds for axis {} (size {})", end, d, shape[d] + ); + new_offset += start * strides[d]; + if !collapse { + new_shape.push(end - start); + new_strides.push(strides[d]); + } + } + + NdArrayV::new(self.source.clone(), new_offset, &new_shape, &new_strides) + } + + // *** Axis manipulation *************************************** + + /// Transposed view with the axis order reversed. Zero-copy - only + /// the shape and stride metadata reorder. A 1D view returns itself + /// unchanged. + pub fn transpose(&self) -> NdArrayV { + let shape: Vec = self.dims.shape().iter().rev().copied().collect(); + let strides: Vec = self.dims.strides().iter().rev().copied().collect(); + NdArrayV::new(self.source.clone(), self.offset, &shape, &strides) + } + + /// View with axes reordered by the given permutation. Zero-copy. + /// + /// `perm[d]` names the source axis that becomes axis `d` of the + /// result. Panics unless `perm` is a permutation of `0..ndim`. + pub fn permute_axes(&self, perm: &[usize]) -> NdArrayV { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + let ndim = shape.len(); + assert_eq!( + perm.len(), ndim, + "permute_axes: expected {} axes, got {}", ndim, perm.len() + ); + let mut seen = vec![false; ndim]; + for &ax in perm { + assert!(ax < ndim, "permute_axes: axis {} out of bounds for {}D view", ax, ndim); + assert!(!seen[ax], "permute_axes: axis {} repeated", ax); + seen[ax] = true; + } + let new_shape: Vec = perm.iter().map(|&ax| shape[ax]).collect(); + let new_strides: Vec = perm.iter().map(|&ax| strides[ax]).collect(); + NdArrayV::new(self.source.clone(), self.offset, &new_shape, &new_strides) + } + + /// View with two axes swapped. Zero-copy. + pub fn swap_axes(&self, a: usize, b: usize) -> NdArrayV { + let ndim = self.dims.ndim(); + assert!( + a < ndim && b < ndim, + "swap_axes: axes ({}, {}) out of bounds for {}D view", a, b, ndim + ); + let mut shape = self.dims.shape().to_vec(); + let mut strides = self.dims.strides().to_vec(); + shape.swap(a, b); + strides.swap(a, b); + NdArrayV::new(self.source.clone(), self.offset, &shape, &strides) + } + + // *** Conversions ********************************************* + + /// Export this view as a legacy DLPack tensor without copying. The + /// window offset carries through DLPack's `byte_offset` field, and + /// the view strides carry as element strides. + #[cfg(feature = "dlpack")] + pub fn to_dlpack(self) -> DLPackTensor { + export_view_to_dlpack(self) + } + + /// Export this view as a DLPack 1.x versioned tensor without + /// copying. The read-only flag is set whenever another reference to + /// the source buffer exists. + #[cfg(feature = "dlpack")] + pub fn to_dlpack_versioned(self) -> DLPackTensorVersioned { + export_view_to_dlpack_versioned(self) + } + + // *** Materialisation ***************************************** + + /// Materialise this view as an owned NdArray. + pub fn to_ndarray(&self) -> NdArray { + let flat: Vec64 = self.into_iter().collect(); + let mut arr = NdArray::from_slice(&flat, self.dims.shape()); + arr.name = self.source.name.clone(); + arr + } + + // *** Parallel iteration (rayon) ****************************** + + /// Parallel iterator over axis-0 observations. Each item is the + /// observation index and a zero-copy `NdArrayV` view. + #[cfg(feature = "parallel_proc")] + pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator)> + '_ + where + T: Send + Sync, + { + use rayon::prelude::*; + assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher view"); + let n_obs = self.dims.shape()[0]; + (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i))) + } + + /// Iterate one logical axis-0 run identified by its flattened outer + /// index. This composes the logical iterator for SuperNdArrayV without + /// materialising its slices. + pub(crate) fn iter_axis0_run(&self, run_idx: usize) -> impl ExactSizeIterator + '_ { + assert!(self.ndim() > 0, "axis-0 iteration requires an axis 0"); + let n_runs: usize = self.shape()[1..].iter().product(); + assert!(run_idx < n_runs, "axis-0 run {} out of bounds ({})", run_idx, n_runs); + + let mut rem = run_idx; + let mut offset = self.offset; + for d in 1..self.ndim() { + offset += (rem % self.shape()[d]) * self.strides()[d]; + rem /= self.shape()[d]; + } + let stride = self.strides()[0]; + (0..self.shape()[0]).map(move |i| self.source.data.as_slice()[offset + i * stride]) + } +} + +/// Materialise a 2D view as a Matrix. +#[cfg(feature = "matrix")] +impl NdArrayV { + /// Materialise a 2D view as a Matrix. + pub fn to_matrix(&self) -> Result { + self.to_ndarray().to_matrix() + } +} + +impl NdArrayV { + /// Apply a function to every logical element, materialising a new + /// compact [`NdArray`] with this view's shape and the parent's name. + pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray { + let flat: Vec64 = self.into_iter().map(f).collect(); + let mut result = NdArray::from_slice(&flat, self.dims.shape()); + result.name = self.source.name.clone(); + result + } +} + +// *** Axis selection: view.s(nd![1..4, 2]) ************************ + +/// Selection across every axis at once, delegating to `slice`. Single +/// indices collapse their dimension, and contiguous ranges keep it. +/// Zero-copy. +#[cfg(feature = "select")] +impl AxisSelection for NdArrayV { + type View = NdArrayV; + + fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + self.slice(selection) + } + + fn get_axis_count(&self) -> usize { + self.ndim() + } +} + +// *** Row selection: view.r(0..10) ******************************** + +/// Axis-0 observation selection over a view. Contiguous ranges narrow +/// the window zero-copy. Index arrays gather the selected observations +/// into an owned array wrapped in a full view. +#[cfg(feature = "select")] +impl RowSelection for NdArrayV { + type View = NdArrayV; + + fn r(&self, selection: S) -> NdArrayV { + assert!(self.ndim() > 0, "row selection requires an axis 0"); + let n_obs = self.dims.shape()[0]; + let indices = selection.resolve_indices(n_obs); + if selection.is_contiguous() { + let start = indices.first().copied().unwrap_or(0); + let ranges: Vec> = std::iter::once(start..start + indices.len()) + .chain(self.dims.shape()[1..].iter().map(|&n| 0..n)) + .collect(); + let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect(); + return self.slice(&refs); + } + NdArrayV::from_ndarray(gather_obs_impl( + &indices, + self.dims.shape(), + self.source.name.clone(), + |idx| self.get(idx), + )) + } + + fn get_row_count(&self) -> usize { + assert!(self.ndim() > 0, "row count requires an axis 0"); + self.dims.shape()[0] + } +} + +// *** IntoIterator ************************************************ + +/// Iterating a view works the same as iterating an NdArray: contiguous +/// runs along axis 0, cache-friendly, no per-element offset arithmetic. +impl<'a, T: Float> IntoIterator for &'a NdArrayV { + type Item = T; + type IntoIter = NdArrayIter<'a, T>; + + fn into_iter(self) -> NdArrayIter<'a, T> { + let shape = self.dims.shape(); + let strides = self.dims.strides(); + if shape.is_empty() { + return NdArrayIter { + buf: self.source.data.as_slice(), + n_inner: 1, + inner_stride: 1, + run_offsets: vec![self.offset], + run_idx: 0, + inner_idx: 0, + total: 1, + yielded: 0, + }; + } + let n_inner = shape[0]; + let n_runs: usize = shape[1..].iter().product(); + + let mut run_offsets = Vec::with_capacity(n_runs); + if shape.len() <= 1 { + run_offsets.push(self.offset); + } else { + let outer_shape = &shape[1..]; + let outer_strides = &strides[1..]; + let mut outer_indices = vec![0usize; outer_shape.len()]; + for _ in 0..n_runs { + let off: usize = self.offset + outer_indices.iter() + .zip(outer_strides.iter()) + .map(|(&i, &s)| i * s) + .sum::(); + run_offsets.push(off); + let mut carry = true; + for d in 0..outer_shape.len() { + if carry { + outer_indices[d] += 1; + if outer_indices[d] < outer_shape[d] { + carry = false; + } else { + outer_indices[d] = 0; + } + } + } + } + } + + NdArrayIter { + buf: self.source.data.as_slice(), + n_inner, + inner_stride: strides[0], + run_offsets, + run_idx: 0, + inner_idx: 0, + total: self.len(), + yielded: 0, + } + } +} + +// *** Trait implementations *************************************** + +impl Shape for NdArrayV { + fn shape(&self) -> ShapeDim { + match self.dims.ndim() { + 0 => ShapeDim::Rank0(1), + 1 => ShapeDim::Rank1(self.dims.shape()[0]), + 2 => ShapeDim::Rank2 { + rows: self.dims.shape()[0], + cols: self.dims.shape()[1], + }, + _ => ShapeDim::RankN(self.dims.shape().to_vec()), + } + } +} + +impl PartialEq for NdArrayV { + fn eq(&self, other: &Self) -> bool { + if self.dims.shape() != other.dims.shape() { return false; } + self.into_iter() + .zip(other.into_iter()) + .all(|(a, b)| a == b) + } +} + +// *** Tuple indexing ********************************************** + +impl Index<()> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (): ()) -> &T { + &self.source.data.as_slice()[self.offset_of(&[])] + } +} + +impl Index<(usize,)> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (i,): (usize,)) -> &T { + &self.source.data.as_slice()[self.offset_of(&[i])] + } +} + +impl Index<(usize, usize)> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (i, j): (usize, usize)) -> &T { + &self.source.data.as_slice()[self.offset_of(&[i, j])] + } +} + +impl Index<(usize, usize, usize)> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (i, j, k): (usize, usize, usize)) -> &T { + &self.source.data.as_slice()[self.offset_of(&[i, j, k])] + } +} + +impl Index<(usize, usize, usize, usize)> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (i, j, k, l): (usize, usize, usize, usize)) -> &T { + &self.source.data.as_slice()[self.offset_of(&[i, j, k, l])] + } +} + +impl Index<(usize, usize, usize, usize, usize)> for NdArrayV { + type Output = T; + #[inline] + fn index(&self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &T { + &self.source.data.as_slice()[self.offset_of(&[i, j, k, l, m])] + } +} + +// *** Debug ******************************************************* + +impl fmt::Debug for NdArrayV { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, "NdArrayV: {:?} [{}D, offset={}]", + self.dims.shape(), self.ndim(), self.offset, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::enums::shape_dim::ShapeDim; + + #[test] + fn from_ndarray() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + assert_eq!(v.shape(), &[3, 2]); + assert_eq!(v.len(), 6); + assert_eq!(v[(0, 0)], 1.0); + assert_eq!(v[(2, 1)], 6.0); + } + + #[test] + fn try_new_validates_and_constructs_view() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]); + let v = NdArrayV::try_new(a, 1, &[2], &[1]).unwrap(); + assert_eq!(v.shape(), &[2]); + assert_eq!(v.get(&[0]), 2.0); + assert_eq!(v.get(&[1]), 3.0); + } + + #[test] + fn try_new_accepts_rank_zero_scalar_view() { + let a = NdArray::from_slice(&[10.0, 20.0], &[2]); + let v = NdArrayV::try_new(a, 1, &[], &[]).unwrap(); + assert!(v.shape().is_empty()); + assert!(v.strides().is_empty()); + assert_eq!(v.len(), 1); + assert_eq!(v[()], 20.0); + assert_eq!((&v).into_iter().collect::>(), vec![20.0]); + assert_eq!(Shape::shape(&v), ShapeDim::Rank0(1)); + } + + #[test] + fn try_new_rejects_invalid_rank_and_span() { + let a = NdArray::::new(&[4]); + assert!(NdArrayV::try_new(a.clone(), 0, &[2, 2], &[1]).is_err()); + assert!(NdArrayV::try_new(a.clone(), 3, &[2], &[1]).is_err()); + assert!(NdArrayV::try_new(a, usize::MAX, &[2], &[1]).is_err()); + } + + #[test] + fn try_new_accepts_empty_view_at_buffer_end() { + let a = NdArray::::new(&[4]); + let v = NdArrayV::try_new(a, 4, &[0], &[1]).unwrap(); + assert!(v.is_empty()); + } + + #[test] + #[should_panic(expected = "NdArrayV::new")] + fn new_panics_on_invalid_span() { + let a = NdArray::::new(&[4]); + let _ = NdArrayV::new(a, 3, &[2], &[1]); + } + + #[test] + fn col_access() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + assert_eq!(v.col(0), &[1.0, 2.0, 3.0]); + assert_eq!(v.col(1), &[4.0, 5.0, 6.0]); + } + + #[test] + fn columns() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + let cols = v.columns(); + assert_eq!(cols.len(), 2); + assert_eq!(cols[0], &[1.0, 2.0, 3.0]); + assert_eq!(cols[1], &[4.0, 5.0, 6.0]); + } + + #[test] + #[cfg(feature = "select")] + fn row_selection_on_view() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = a.as_view(); + // Contiguous narrows zero-copy. + let sub = v.r(1..3); + assert_eq!(sub.shape(), &[2, 2]); + assert_eq!(sub.get(&[0, 1]), 5.0); + // Index arrays gather in order, and the source is unaffected. + let picked = v.r(&[2, 0]); + assert_eq!(picked.get(&[0, 0]), 3.0); + assert_eq!(picked.get(&[1, 1]), 4.0); + } + + #[test] + fn apply_on_view() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]); + let out = a.as_view().apply(|x| x + 1.0); + assert_eq!(out.shape(), &[2, 2]); + assert_eq!(out.get(&[1, 1]), 5.0); + } + + #[test] + fn obs_access() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + let obs0: Vec = (&v.obs(0)).into_iter().collect(); + let obs2: Vec = (&v.obs(2)).into_iter().collect(); + assert_eq!(obs0, vec![1.0, 4.0]); + assert_eq!(obs2, vec![3.0, 6.0]); + } + + #[test] + fn iteration() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + #[test] + fn iteration_2d() { + let data: Vec = (1..=20).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[10, 2]); + let v = NdArrayV::from_ndarray(a); + let vals: Vec = (&v).into_iter().collect(); + assert_eq!(vals.len(), 20); + assert_eq!(&vals[..10], &data[..10]); + assert_eq!(&vals[10..], &data[10..]); + } + + #[test] + fn with_offset() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let stride = a.strides()[1]; + // View into column 1 only as a 1D view + let v = NdArrayV::new(a.clone(), stride, &[3], &[1]); + assert_eq!(v.shape(), &[3]); + assert_eq!(v[(0,)], 4.0); + assert_eq!(v[(1,)], 5.0); + assert_eq!(v[(2,)], 6.0); + } + + #[test] + fn to_ndarray() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + let b = v.to_ndarray(); + assert_eq!(b.shape(), &[3, 2]); + assert_eq!(b.col(0), &[1.0, 2.0, 3.0]); + } + + #[test] + fn transpose_2d_view() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let t = NdArrayV::from_ndarray(a).transpose(); + assert_eq!(t.shape(), &[2, 3]); + assert_eq!(t[(0, 0)], 1.0); + assert_eq!(t[(1, 0)], 4.0); + assert_eq!(t[(0, 2)], 3.0); + assert_eq!(t[(1, 2)], 6.0); + // Materialising the transposed view matches the owned transpose. + let owned = t.to_ndarray(); + assert_eq!(owned.shape(), &[2, 3]); + assert_eq!(owned.get(&[1, 1]), 5.0); + } + + #[test] + fn transpose_3d_view_matches_materialised() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let t = a.as_view().transpose(); + assert_eq!(t.shape(), &[4, 3, 2]); + let materialised = a.transpose(); + for i in 0..4 { + for j in 0..3 { + for k in 0..2 { + assert_eq!(t[(i, j, k)], materialised.get(&[i, j, k])); + } + } + } + } + + #[test] + fn permute_axes_view() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let p = a.as_view().permute_axes(&[2, 0, 1]); + assert_eq!(p.shape(), &[4, 2, 3]); + for i in 0..2 { + for j in 0..3 { + for k in 0..4 { + assert_eq!(p[(k, i, j)], a.get(&[i, j, k])); + } + } + } + } + + #[test] + #[should_panic(expected = "permute_axes")] + fn permute_axes_rejects_repeat() { + let a = NdArray::::new(&[2, 3, 4]); + let _ = a.as_view().permute_axes(&[0, 0, 1]); + } + + #[test] + fn swap_axes_view() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let s = a.as_view().swap_axes(0, 2); + assert_eq!(s.shape(), &[4, 3, 2]); + assert_eq!(s[(3, 2, 1)], a.get(&[1, 2, 3])); + assert_eq!(s[(0, 0, 0)], a.get(&[0, 0, 0])); + } + + #[test] + #[should_panic(expected = "obs()")] + fn obs_on_1d_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let _ = a.as_view().obs(1); + } + + #[cfg(feature = "matrix")] + #[test] + fn to_matrix() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let v = NdArrayV::from_ndarray(a); + let mat = v.to_matrix().unwrap(); + assert_eq!(mat.n_rows, 3); + assert_eq!(mat.n_cols, 2); + } + + #[test] + fn blas_params() { + let a = NdArray::::new(&[10, 5]); + let v = NdArrayV::from_ndarray(a); + assert_eq!(v.m(), 10); + assert_eq!(v.n(), 5); + assert_eq!(v.lda(), 10); + } + + #[test] + fn eq() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]); + let v1 = NdArrayV::from_ndarray(a.clone()); + let v2 = NdArrayV::from_ndarray(a); + assert_eq!(v1, v2); + } + + #[test] + fn shape_trait() { + let a = NdArray::::new(&[3, 4]); + let v = NdArrayV::from_ndarray(a); + assert_eq!(Shape::shape(&v), ShapeDim::Rank2 { rows: 3, cols: 4 }); + } + + #[test] + #[should_panic(expected = "unit stride")] + fn col_on_transposed_view_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let t = a.as_view().transpose(); + let _ = t.col(0); + } + + #[test] + #[should_panic(expected = "unit stride")] + fn columns_on_transposed_view_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let t = a.as_view().transpose(); + let _ = t.columns(); + } + + #[test] + #[should_panic(expected = "unit stride")] + fn lda_on_transposed_view_panics() { + let a = NdArray::::new(&[3, 2]); + let t = a.as_view().transpose(); + let _ = t.lda(); + } + + #[test] + #[should_panic(expected = "index 3 out of bounds for axis 0")] + fn obs_out_of_bounds_panics() { + let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let _ = a.as_view().obs(3); + } + + #[test] + #[should_panic(expected = "swap_axes: axes (0, 3) out of bounds")] + fn swap_axes_out_of_bounds_panics() { + let a = NdArray::::new(&[2, 3, 4]); + let _ = a.as_view().swap_axes(0, 3); + } + + #[test] + #[should_panic(expected = "permute_axes: axis 3 out of bounds")] + fn permute_axes_out_of_bounds_axis_panics() { + let a = NdArray::::new(&[2, 3, 4]); + let _ = a.as_view().permute_axes(&[0, 1, 3]); + } + + #[test] + fn get_unchecked_matches_get() { + let data: Vec = (1..=24).map(|x| x as f64).collect(); + let a = NdArray::from_slice(&data, &[2, 3, 4]); + let t = a.as_view().transpose(); + for i in 0..4 { + for j in 0..3 { + for k in 0..2 { + let idx = [i, j, k]; + // SAFETY: indices are within shape + assert_eq!(unsafe { t.get_unchecked(&idx) }, t.get(&idx)); + } + } + } + } +} diff --git a/src/structs/xarray.rs b/src/structs/xarray.rs new file mode 100644 index 0000000..b8af9e1 --- /dev/null +++ b/src/structs/xarray.rs @@ -0,0 +1,1870 @@ +//! # `XArray` +//! +//! A labelled N-dimensional array with named axes and optional coordinates. +//! +//! `XArray` associates an [`NdArray`] with dimension names and, optionally, +//! coordinate labels for each axis. It supports both positional selection and +//! coordinate-based queries without changing the underlying numerical model. +//! +//! Where `NdArray` addresses values by position, such as `[3, 7]`, `XArray` +//! can address them using domain values such as latitude `51.5`, ticker `"NQ"`, +//! or timestamps within a specified interval. +//! +//! This representation is well-suited to spatial, climate, sensor, time-series, and +//! other datasets whose dimensions carry semantic meaning. +//! +//! ## Storage +//! +//! An `XArray` contains either an owned [`NdArray`] or a zero-copy [`NdArrayV`] +//! view behind the same public interface. Selections can therefore remain +//! labelled `XArray` values without materialising or copying the selected data. +//! +//! ## Usage +//! +//! `ignore +//! let xa = XArray::new(data, &["observation", "feature"]); +//! +//! xa.ax("feature"); // Return axis metadata. +//! xa.dim("feature"); // Resolve the axis position: 1. +//! xa.at("feature", 2.0); // Select a numeric coordinate. +//! xa.at("ticker", "NQ"); // Select a string coordinate. +//! xa.at("time", 1_700_000_000_000i64); // Select a datetime tick. +//! xa.nearest("time", ts); // Select the closest coordinate. +//! xa.between("observation", 0.0, 50.0); // Select a coordinate interval. +//! xa.select(&[("lat", &(0..3))]); // Select by positional range. +//! ` + +use std::fmt; + +use crate::enums::error::MinarrowError; +use crate::enums::shape_dim::ShapeDim; +use crate::ffi::arrow_dtype::ArrowType; +use crate::structs::ndarray::NdArray; +#[cfg(all(feature = "views", feature = "select"))] +use crate::traits::selection::{AxisSelection, DataSelector}; +#[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] +use crate::Scalar; +#[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] +use crate::{NumericArray, TextArray}; +#[cfg(all( + feature = "views", + feature = "select", + feature = "scalar_type", + feature = "datetime" +))] +use crate::TemporalArray; +#[cfg(all(feature = "views", feature = "select"))] +use std::ops::Range; +use crate::traits::type_unions::Float; +use crate::traits::{concatenate::Concatenate, shape::Shape}; +use crate::{Array, Field, StringArray, Table}; + +#[cfg(feature = "views")] +use crate::structs::views::ndarray_view::NdArrayV; + +// **************************************************************** +// Dispatch macros +// **************************************************************** + +/// Dispatch to the inner NdArray/NdArrayV. +#[cfg(feature = "views")] +macro_rules! delegate { + ($self:expr, $method:ident ( $($arg:expr),* )) => { + match &$self.data { + NdArrayE::Owned(nd) => nd.$method($($arg),*), + NdArrayE::View(v) => v.$method($($arg),*), + } + }; +} + +#[cfg(not(feature = "views"))] +macro_rules! delegate { + ($self:expr, $method:ident ( $($arg:expr),* )) => { + match &$self.data { + NdArrayE::Owned(nd) => nd.$method($($arg),*), + } + }; +} + +// **************************************************************** +// XArray +// **************************************************************** + +/// Labelled N-dimensional array with named dimensions and coordinate-based indexing. +/// +/// XArray wraps an [`NdArray`] or [`NdArrayV`] with per-axis names and optional +/// coordinate labels, enabling selection by value rather than raw position. +/// Owned data and selection views expose the same container interface. +/// +/// # Construction +/// ``` +/// use minarrow::structs::ndarray::NdArray; +/// use minarrow::structs::xarray::{XArray, Axis}; +/// +/// // Name the dimensions, no coordinate labels +/// let data = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); +/// let xa = XArray::new(data, &["station", "measurement"]); +/// assert_eq!(xa.dim("measurement"), 1); +/// ``` +/// +/// # Coordinate-based selection +/// Assign coordinate labels to an axis, then select by value: +/// ``` +/// # use minarrow::structs::ndarray::NdArray; +/// # use minarrow::structs::xarray::XArray; +/// # use minarrow::arr_f64; +/// // 5 stations, 2 measurements each +/// let data = NdArray::from_slice( +/// &[10.0, 20.0, 30.0, 40.0, 50.0, 1.0, 2.0, 3.0, 4.0, 5.0], +/// &[5, 2], +/// ); +/// let mut xa = XArray::new(data, &["station", "measurement"]); +/// +/// // Label the station axis with latitude values +/// xa.assign_coords("station", arr_f64![-33.8, 35.7, 51.5, 40.7, -22.9]); +/// +/// // Select stations between latitudes 35 and 52 +/// let subset = xa.between("station", 35.0, 52.0); +/// assert_eq!(subset.shape(), vec![3, 2]); // 3 stations matched +/// ``` +/// +/// # Positional selection +/// Select by axis name and index/range without coordinates: +/// ``` +/// # use minarrow::structs::ndarray::NdArray; +/// # use minarrow::structs::xarray::XArray; +/// let data = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); +/// let xa = XArray::new(data, &["obs", "feat"]); +/// +/// // Range on one axis - keeps the dimension +/// let sub = xa.select(&[("obs", &(0..2))]); +/// assert_eq!(sub.shape(), vec![2, 2]); +/// +/// // Single index collapses the dimension +/// let col0 = xa.select(&[("feat", &0)]); +/// assert_eq!(col0.ndim(), 1); +/// ``` +#[derive(Clone)] +pub struct XArray { + data: NdArrayE, + axes: Vec, +} + +impl XArray { + /// Create with named dimensions, no coordinates. + pub fn new(data: NdArray, dim_names: &[&str]) -> Self { + assert_eq!( + data.ndim(), dim_names.len(), + "XArray: {} dim names for {}D array", dim_names.len(), data.ndim() + ); + let axes = dim_names.iter().map(|n| Axis::named(*n)).collect(); + XArray { data: NdArrayE::Owned(data), axes } + } + + /// Create with fully specified axes. + pub fn with_axes(data: NdArray, axes: Vec) -> Self { + assert_eq!( + data.ndim(), axes.len(), + "XArray: {} axes for {}D array", axes.len(), data.ndim() + ); + for (i, ax) in axes.iter().enumerate() { + if let Some(ref coords) = ax.coords { + assert_eq!( + coords.len(), data.shape()[i], + "XArray: axis '{}' has {} coords but dimension size is {}", + ax.name, coords.len(), data.shape()[i] + ); + } + } + XArray { data: NdArrayE::Owned(data), axes } + } + + /// Create from NdArray with auto-generated dim names. + pub fn from_ndarray(data: NdArray) -> Self { + let axes = (0..data.ndim()) + .map(|i| Axis::named(format!("dim_{}", i))) + .collect(); + XArray { data: NdArrayE::Owned(data), axes } + } + + /// Wrap an NdArrayV view with axes (zero-copy). + #[cfg(feature = "views")] + pub fn from_view(view: NdArrayV, axes: Vec) -> Self { + assert_eq!( + view.ndim(), axes.len(), + "XArray: {} axes for {}D view", axes.len(), view.ndim() + ); + for (i, ax) in axes.iter().enumerate() { + if let Some(ref coords) = ax.coords { + assert_eq!( + coords.len(), view.shape()[i], + "XArray: axis '{}' has {} coords but dimension size is {}", + ax.name, coords.len(), view.shape()[i] + ); + } + } + XArray { data: NdArrayE::View(view), axes } + } + + // **************************************************************** + // Axis access + // **************************************************************** + + /// Get axis by name. Returns None if not found. + pub fn try_ax(&self, name: &str) -> Option<&Axis> { + self.axes.iter().find(|a| a.name == name) + } + + /// Get axis by name. Panics if not found. + pub fn ax(&self, name: &str) -> &Axis { + self.try_ax(name) + .unwrap_or_else(|| panic!("XArray: no axis named '{}'", name)) + } + + /// Get the position of a named axis. Returns None if not found. + pub fn try_dim(&self, name: &str) -> Option { + self.axes.iter().position(|a| a.name == name) + } + + /// Get the position of a named axis. Panics if not found. + pub fn dim(&self, name: &str) -> usize { + self.try_dim(name) + .unwrap_or_else(|| panic!("XArray: no axis named '{}'", name)) + } + + /// All axes. + #[inline] + pub fn axes(&self) -> &[Axis] { &self.axes } + + /// Dim names. + pub fn dim_names(&self) -> Vec<&str> { + self.axes.iter().map(|a| a.name.as_str()).collect() + } + + // **************************************************************** + // Delegated data access + // **************************************************************** + + /// Consume and return the inner NdArray, materialising if it is a view. + pub fn into_ndarray(self) -> NdArray { + match self.data { + NdArrayE::Owned(nd) => nd, + #[cfg(feature = "views")] + NdArrayE::View(v) => v.to_ndarray(), + } + } + + /// Materialise to an owned NdArray if currently a view. + pub fn to_owned(&self) -> XArray { + match &self.data { + NdArrayE::Owned(_) => self.clone(), + #[cfg(feature = "views")] + NdArrayE::View(v) => XArray { + data: NdArrayE::Owned(v.to_ndarray()), + axes: self.axes.clone(), + }, + } + } + + /// Borrow the data as a zero-copy NdArrayV, regardless of whether this + /// XArray currently owns its NdArray or already holds a view. + #[cfg(feature = "views")] + pub fn as_view(&self) -> NdArrayV { + match &self.data { + NdArrayE::Owned(nd) => nd.as_view(), + NdArrayE::View(v) => v.clone(), + } + } + + /// True if backed by a single owned NdArray. + pub fn is_owned(&self) -> bool { + matches!(&self.data, NdArrayE::Owned(_)) + } + + /// Borrow the inner storage for crate-internal dispatch. + #[cfg(feature = "size")] + #[inline] + pub(crate) fn storage(&self) -> &NdArrayE { + &self.data + } + + #[inline] + pub fn ndim(&self) -> usize { delegate!(self, ndim()) } + + pub fn shape(&self) -> Vec { + match &self.data { + NdArrayE::Owned(nd) => nd.shape().to_vec(), + #[cfg(feature = "views")] + NdArrayE::View(v) => v.shape().to_vec(), + } + } + + #[inline] + pub fn strides(&self) -> &[usize] { delegate!(self, strides()) } + + #[inline] + pub fn len(&self) -> usize { delegate!(self, len()) } + + #[inline] + pub fn is_empty(&self) -> bool { delegate!(self, is_empty()) } + + /// Zero-copy view of a single observation (axis-0 element). + /// + /// Returns an (N-1)-dimensional `NdArrayV` view regardless of + /// the inner storage mode. + #[cfg(feature = "views")] + pub fn obs(&self, idx: usize) -> NdArrayV { + match &self.data { + NdArrayE::Owned(nd) => nd.as_view().obs(idx), + #[cfg(feature = "views")] + NdArrayE::View(v) => v.obs(idx), + } + } + + #[inline] + pub fn m(&self) -> i32 { delegate!(self, m()) } + + #[inline] + pub fn n(&self) -> i32 { delegate!(self, n()) } + + #[inline] + pub fn lda(&self) -> i32 { delegate!(self, lda()) } + + /// Single element access. + pub fn get(&self, indices: &[usize]) -> T { delegate!(self, get(indices)) } + + /// Mutable element access. Triggers copy-on-write if the owned array is + /// shared with views; an XArray backed by a view cannot be mutated. + pub fn set(&mut self, indices: &[usize], value: T) { + match &mut self.data { + NdArrayE::Owned(nd) => nd.set(indices, value), + #[cfg(feature = "views")] + NdArrayE::View(_) => panic!("XArray: cannot mutate a view"), + } + } + + /// Parallel iterator over axis-0 observations. Each item is the + /// observation index and a zero-copy `NdArrayV` view. + #[cfg(all(feature = "parallel_proc", feature = "views"))] + pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator)> + '_ + where + T: Send + Sync, + { + use rayon::prelude::*; + assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher array"); + let n_obs = self.shape()[0]; + (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i))) + } + + // **************************************************************** + // Apply + // **************************************************************** + + /// Apply a function to every logical element, returning a new labelled + /// array with the same axes. View storage materialises to owned. + pub fn apply(&self, f: impl Fn(T) -> T) -> XArray { + let data = match &self.data { + NdArrayE::Owned(nd) => NdArrayE::Owned(nd.apply(f)), + #[cfg(feature = "views")] + NdArrayE::View(v) => NdArrayE::Owned(v.apply(f)), + }; + XArray { data, axes: self.axes.clone() } + } + + /// Apply a function to every logical element in place. An XArray backed + /// by a view cannot be mutated, matching `set`. + pub fn apply_mut(&mut self, f: impl Fn(T) -> T) { + match &mut self.data { + NdArrayE::Owned(nd) => nd.apply_mut(f), + #[cfg(feature = "views")] + NdArrayE::View(_) => panic!("XArray: cannot mutate a view"), + } + } + + // **************************************************************** + // Label management + // **************************************************************** + + /// Rename an axis. + pub fn rename_dim(&mut self, old: &str, new: &str) { + let idx = self.dim(old); + self.axes[idx].name = new.to_string(); + } + + /// Assign or replace coordinates for a named axis. + pub fn assign_coords(&mut self, dim_name: &str, coords: Array) { + let idx = self.dim(dim_name); + let dim_size = self.shape()[idx]; + assert_eq!( + coords.len(), dim_size, + "XArray: coords length {} does not match axis '{}' size {}", + coords.len(), dim_name, dim_size + ); + self.axes[idx].coords = Some(coords); + } + + /// Remove coordinates from a named axis. + pub fn drop_coords(&mut self, dim_name: &str) { + let idx = self.dim(dim_name); + self.axes[idx].coords = None; + } + + // **************************************************************** + // Positional selection: .select() + // **************************************************************** + + /// Select sub-arrays by named axis positions. Returns zero-copy XArray + /// backed by NdArrayV. + /// + /// Single indices collapse that dimension. Ranges keep it. + /// + /// # Examples + /// ```ignore + /// xa.select(&[("lat", &(0..3))]) // single axis range + /// xa.select(&[("lat", &(0..3)), ("lon", &2)]) // multi-axis mixed + /// ``` + #[cfg(all(feature = "views", feature = "select"))] + pub fn select(&self, selection: &[(&str, &dyn DataSelector)]) -> XArray { + let shape = self.shape(); + + // Default: full range on every axis + let full_ranges: Vec> = shape.iter().map(|&n| 0..n).collect(); + let mut refs: Vec<&dyn DataSelector> = full_ranges.iter().map(|r| r as _).collect(); + let mut new_axes: Vec> = + self.axes.iter().cloned().map(Some).collect(); + + // Apply named selections. Collapsed axes drop their labels, and + // windowed axes carry their coordinates through, narrowed. + for (name, sel) in selection { + let idx = self.dim(name); + refs[idx] = *sel; + let (start, end, collapse) = sel.resolve_axis(shape[idx]); + if collapse { + new_axes[idx] = None; + } else { + let mut narrowed = Axis::named(&self.axes[idx].name); + if let Some(ref coords) = self.axes[idx].coords { + narrowed.coords = Some(coords.slice_clone(start, end - start)); + } + new_axes[idx] = Some(narrowed); + } + } + + let view = self.slice(&refs); + let result_axes: Vec = new_axes.into_iter().flatten().collect(); + XArray { data: NdArrayE::View(view), axes: result_axes } + } + + /// Slice the underlying NdArray/NdArrayV positionally and zero-copy. + /// For named axis selection, use `.select()` instead. + #[cfg(all(feature = "views", feature = "select"))] + pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + match &self.data { + NdArrayE::Owned(nd) => nd.slice(selection), + NdArrayE::View(v) => v.slice(selection), + } + } + + // **************************************************************** + // Coordinate value selection: .at() and .between() + // **************************************************************** + + /// Select a single position on a named axis by coordinate value. + /// Accepts numeric, string, and datetime values. Collapses that + /// dimension. Returns an error if the value is not found. Float + /// coordinates match by IEEE equality, so NaN never matches and + /// derived values may miss - `nearest` tolerates rounding. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn try_at(&self, dim_name: &str, value: impl Into) -> Result, MinarrowError> { + let dim_idx = self.dim(dim_name); + let pos = self.axes[dim_idx].try_coord_pos(&value.into())?; + + let shape = self.shape(); + let full_ranges: Vec> = shape.iter().map(|&n| 0..n).collect(); + let mut refs: Vec<&dyn DataSelector> = full_ranges.iter().map(|r| r as _).collect(); + refs[dim_idx] = &pos; + let new_axes: Vec = self + .axes + .iter() + .enumerate() + .filter(|(i, _)| *i != dim_idx) + .map(|(_, ax)| ax.clone()) + .collect(); + + let view = self.slice(&refs); + Ok(XArray { data: NdArrayE::View(view), axes: new_axes }) + } + + /// Select a single position by coordinate value. Panics if not found. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn at(&self, dim_name: &str, value: impl Into) -> XArray { + self.try_at(dim_name, value) + .unwrap_or_else(|e| panic!("{}", e)) + } + + /// Select a range by coordinate value bounds (inclusive). + /// Accepts numeric, string, and datetime bounds. Returns an error + /// if no values fall in the range, or if the matching coordinates + /// do not form a contiguous run i.e. the axis is not monotonic over + /// the requested bounds - sort the axis or gather by position for + /// unsorted coordinates. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn try_between( + &self, + dim_name: &str, + low: impl Into, + high: impl Into, + ) -> Result, MinarrowError> { + let dim_idx = self.dim(dim_name); + let (start, end) = self.axes[dim_idx].try_coord_range(&low.into(), &high.into())?; + + let shape = self.shape(); + let full_ranges: Vec> = shape.iter().map(|&n| 0..n).collect(); + let mut refs: Vec<&dyn DataSelector> = full_ranges.iter().map(|r| r as _).collect(); + let window = start..end; + refs[dim_idx] = &window; + let new_axes: Vec = self + .axes + .iter() + .enumerate() + .map(|(i, ax)| { + if i == dim_idx { + let mut narrowed = Axis::named(&ax.name); + if let Some(ref coords) = ax.coords { + narrowed.coords = Some(coords.slice_clone(start, end - start)); + } + narrowed + } else { + ax.clone() + } + }) + .collect(); + + let view = self.slice(&refs); + Ok(XArray { data: NdArrayE::View(view), axes: new_axes }) + } + + /// Select a range by coordinate value bounds. Panics if no values + /// match, or if the axis is not monotonic over the requested bounds. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn between( + &self, + dim_name: &str, + low: impl Into, + high: impl Into, + ) -> XArray { + self.try_between(dim_name, low, high) + .unwrap_or_else(|e| panic!("{}", e)) + } + + /// Select the position whose coordinate is closest to `value` on a + /// named axis. Collapses that dimension. Numeric and datetime + /// coordinates only, since text has no distance metric. Returns an + /// error if the axis has no comparable coordinates. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn try_nearest(&self, dim_name: &str, value: impl Into) -> Result, MinarrowError> { + let dim_idx = self.dim(dim_name); + let pos = self.axes[dim_idx].try_coord_nearest(&value.into())?; + + let shape = self.shape(); + let full_ranges: Vec> = shape.iter().map(|&n| 0..n).collect(); + let mut refs: Vec<&dyn DataSelector> = full_ranges.iter().map(|r| r as _).collect(); + refs[dim_idx] = &pos; + let new_axes: Vec = self + .axes + .iter() + .enumerate() + .filter(|(i, _)| *i != dim_idx) + .map(|(_, ax)| ax.clone()) + .collect(); + + let view = self.slice(&refs); + Ok(XArray { data: NdArrayE::View(view), axes: new_axes }) + } + + /// Select the position whose coordinate is closest to `value`. + /// Panics if the axis has no comparable coordinates. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub fn nearest(&self, dim_name: &str, value: impl Into) -> XArray { + self.try_nearest(dim_name, value) + .unwrap_or_else(|e| panic!("{}", e)) + } + + // **************************************************************** + // Axis operations + // **************************************************************** + + /// Transpose (2D only). Reorders axes by name, so the result's + /// dimensions arrive in `dim_order`. Passing the current order + /// returns the array unchanged. + pub fn transpose(&self, dim_order: &[&str]) -> Result, MinarrowError> { + if self.ndim() != 2 { + return Err(MinarrowError::ShapeError { + message: format!("transpose requires 2D, got {}D", self.ndim()), + }); + } + if dim_order.len() != 2 { + return Err(MinarrowError::ShapeError { + message: "transpose: expected 2 dim names".to_string(), + }); + } + if dim_order[0] == dim_order[1] { + return Err(MinarrowError::ShapeError { + message: format!( + "transpose: dim names must be distinct, got '{}' twice", dim_order[0] + ), + }); + } + for name in dim_order { + if self.try_dim(name).is_none() { + return Err(MinarrowError::ShapeError { + message: format!( + "transpose: no axis named '{}', axes are {:?}", + name, + self.dim_names() + ), + }); + } + } + // The current order is a no-op, so the data only transposes when + // the axes actually swap. + if dim_order[0] == self.axes[0].name { + return Ok(self.clone()); + } + let inner = match &self.data { + NdArrayE::Owned(nd) => nd.transpose(), + #[cfg(feature = "views")] + NdArrayE::View(v) => v.transpose().to_ndarray(), + }; + let new_axes = dim_order.iter() + .map(|name| self.ax(name).clone()) + .collect(); + Ok(XArray { data: NdArrayE::Owned(inner), axes: new_axes }) + } + +} + +impl XArray { + /// Convert a 2D XArray to a Table. Uses axis 1 coords as column names + /// if available, otherwise generates names from the dim name. + pub fn to_table(self) -> Result { + if self.ndim() != 2 { + return Err(MinarrowError::ShapeError { + message: format!("to_table requires 2D, got {}D", self.ndim()), + }); + } + let n_cols = self.shape()[1]; + let fields: Vec = if let Some(ref coords) = self.axes[1].coords { + (0..n_cols).map(|i| { + let name = coords.value_to_string(i); + Field::new(name, ArrowType::Float64, false, None) + }).collect() + } else { + (0..n_cols).map(|i| { + Field::new( + format!("{}_{}", self.axes[1].name, i), + ArrowType::Float64, false, None, + ) + }).collect() + }; + self.into_ndarray().to_table(Some(fields)) + } +} + +// **************************************************************** +// NdArrayEnum - owned or view storage +// **************************************************************** + +/// Either owned NdArray or zero-copy NdArrayV. +/// Enables XArray to be a single type regardless of ownership. +#[derive(Clone)] +pub(crate) enum NdArrayE { + /// The array's own shared buffer makes clones a refcount bump and + /// lets views borrow the parent zero-copy, with copy-on-write + /// mutation. + Owned(NdArray), + #[cfg(feature = "views")] + View(NdArrayV), +} + +// **************************************************************** +// Axis +// **************************************************************** + +/// A named dimension with optional coordinate labels. +/// +/// The coords array, when present, must have the same length as the +/// corresponding NdArray dimension. Coordinates may be stored as any +/// Minarrow Array type. Value-based selection resolves numeric, string, +/// and datetime coordinates. +#[derive(Clone, Debug, PartialEq)] +pub struct Axis { + pub name: String, + pub coords: Option, +} + +/// Scans a typed coordinate slice, widening `start`/`end` over positions +/// within the inclusive `[lo, hi]` window and counting the matches, so the +/// caller can detect a span polluted by out-of-range coordinates. +#[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] +macro_rules! coord_window { + ($data:expr, $lo:expr, $hi:expr, $domain:ty, $start:ident, $end:ident, $count:ident) => { + for (i, &v) in $data.iter().enumerate() { + if (v as $domain) >= $lo && (v as $domain) <= $hi { + if i < $start { $start = i; } + if i + 1 > $end { $end = i + 1; } + $count += 1; + } + } + }; +} + +/// Scans string-valued coordinates, widening `start`/`end` over positions +/// within the inclusive lexicographic `[lo, hi]` window and counting the +/// matches, so the caller can detect a span polluted by out-of-range +/// coordinates. +#[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] +macro_rules! coord_window_str { + ($arr:expr, $n:expr, $lo:expr, $hi:expr, $start:ident, $end:ident, $count:ident) => { + for i in 0..$n { + if let Some(s) = $arr.get_str(i) { + if s >= $lo && s <= $hi { + if i < $start { $start = i; } + if i + 1 > $end { $end = i + 1; } + $count += 1; + } + } + } + }; +} + +/// Scans a typed coordinate slice for the position with the smallest +/// distance to the target. Ties resolve to the earliest position. +#[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] +macro_rules! coord_nearest { + ($data:expr, |$v:ident| $dist:expr) => {{ + let mut best = None; + for (i, &$v) in $data.iter().enumerate() { + let dist = $dist; + if best.is_none_or(|(_, bd)| dist < bd) { + best = Some((i, dist)); + } + } + best.map(|(i, _)| i) + }}; +} + +impl Axis { + /// Named axis without coordinates. + pub fn named(name: impl Into) -> Self { + Axis { name: name.into(), coords: None } + } + + /// Named axis with coordinate labels. + pub fn with_coords(name: impl Into, coords: Array) -> Self { + Axis { name: name.into(), coords: Some(coords) } + } + + /// Position of a coordinate value on this axis. String and + /// categorical coordinates match by string equality, datetime + /// coordinates by tick value in the axis's time unit, and numeric + /// coordinates in their native domain (i64, u64, or f64). + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub(crate) fn try_coord_pos(&self, value: &Scalar) -> Result { + let coords = self.coords.as_ref().ok_or_else(|| MinarrowError::ShapeError { + message: format!("axis '{}' has no coordinates for value lookup", self.name), + })?; + let n = coords.len(); + let found = match coords { + Array::TextArray(text) => { + let target = value.try_str().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "string", + message: Some(format!("axis '{}' holds text coordinates", self.name)), + })?; + let target = target.as_str(); + match text { + TextArray::String32(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + #[cfg(feature = "large_string")] + TextArray::String64(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + #[cfg(feature = "default_categorical_8")] + TextArray::Categorical8(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + #[cfg(feature = "extended_categorical")] + TextArray::Categorical16(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + #[cfg(any( + not(feature = "default_categorical_8"), + feature = "extended_categorical" + ))] + TextArray::Categorical32(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + #[cfg(feature = "extended_categorical")] + TextArray::Categorical64(a) => (0..n).find(|&i| a.get_str(i) == Some(target)), + TextArray::Null => None, + } + } + Array::NumericArray(num) => match num { + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int8(a) => + value.try_i64().and_then(|t| a.data.iter().position(|&v| v as i64 == t)), + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int16(a) => + value.try_i64().and_then(|t| a.data.iter().position(|&v| v as i64 == t)), + NumericArray::Int32(a) => + value.try_i64().and_then(|t| a.data.iter().position(|&v| v as i64 == t)), + NumericArray::Int64(a) => + value.try_i64().and_then(|t| a.data.iter().position(|&v| v == t)), + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt8(a) => + value.try_u64().and_then(|t| a.data.iter().position(|&v| v as u64 == t)), + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(a) => + value.try_u64().and_then(|t| a.data.iter().position(|&v| v as u64 == t)), + NumericArray::UInt32(a) => + value.try_u64().and_then(|t| a.data.iter().position(|&v| v as u64 == t)), + NumericArray::UInt64(a) => + value.try_u64().and_then(|t| a.data.iter().position(|&v| v == t)), + NumericArray::Float32(a) => + value.try_f64().and_then(|t| a.data.iter().position(|&v| v as f64 == t)), + NumericArray::Float64(a) => + value.try_f64().and_then(|t| a.data.iter().position(|&v| v == t)), + NumericArray::Null => None, + }, + #[cfg(feature = "datetime")] + Array::TemporalArray(temporal) => { + let target = value.try_i64().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "datetime ticks", + message: Some(format!("axis '{}' holds datetime coordinates", self.name)), + })?; + match temporal { + TemporalArray::Datetime32(a) => a.data.iter().position(|&t| t as i64 == target), + TemporalArray::Datetime64(a) => a.data.iter().position(|&t| t == target), + TemporalArray::Null => None, + } + } + _ => return Err(MinarrowError::TypeError { + from: "coordinate array", to: "comparable value", + message: Some(format!( + "axis '{}' coordinates do not support value lookup", self.name + )), + }), + }; + found.ok_or_else(|| MinarrowError::IndexError(format!( + "value {:?} not found on axis '{}'", value, self.name + ))) + } + + /// Start and end positions covering all coordinates within the + /// inclusive `[low, high]` bounds. String and categorical coordinates + /// compare lexicographically, datetime coordinates by tick value, and + /// numeric coordinates in their native domain (i64, u64, or f64). + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub(crate) fn try_coord_range( + &self, + low: &Scalar, + high: &Scalar, + ) -> Result<(usize, usize), MinarrowError> { + let coords = self.coords.as_ref().ok_or_else(|| MinarrowError::ShapeError { + message: format!("axis '{}' has no coordinates for range lookup", self.name), + })?; + let n = coords.len(); + let mut start = n; + let mut end = 0; + let mut count = 0usize; + match coords { + Array::TextArray(text) => { + let lo = low.try_str().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "string", + message: Some(format!("axis '{}' holds text coordinates", self.name)), + })?; + let hi = high.try_str().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "string", + message: Some(format!("axis '{}' holds text coordinates", self.name)), + })?; + let (lo, hi) = (lo.as_str(), hi.as_str()); + match text { + TextArray::String32(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + #[cfg(feature = "large_string")] + TextArray::String64(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + #[cfg(feature = "default_categorical_8")] + TextArray::Categorical8(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + #[cfg(feature = "extended_categorical")] + TextArray::Categorical16(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + #[cfg(any( + not(feature = "default_categorical_8"), + feature = "extended_categorical" + ))] + TextArray::Categorical32(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + #[cfg(feature = "extended_categorical")] + TextArray::Categorical64(a) => { coord_window_str!(a, n, lo, hi, start, end, count); } + TextArray::Null => {} + } + } + Array::NumericArray(num) => match num { + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int8(a) => if let (Some(lo), Some(hi)) = (low.try_i64(), high.try_i64()) { + coord_window!(a.data, lo, hi, i64, start, end, count); + }, + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int16(a) => if let (Some(lo), Some(hi)) = (low.try_i64(), high.try_i64()) { + coord_window!(a.data, lo, hi, i64, start, end, count); + }, + NumericArray::Int32(a) => if let (Some(lo), Some(hi)) = (low.try_i64(), high.try_i64()) { + coord_window!(a.data, lo, hi, i64, start, end, count); + }, + NumericArray::Int64(a) => if let (Some(lo), Some(hi)) = (low.try_i64(), high.try_i64()) { + coord_window!(a.data, lo, hi, i64, start, end, count); + }, + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt8(a) => if let (Some(lo), Some(hi)) = (low.try_u64(), high.try_u64()) { + coord_window!(a.data, lo, hi, u64, start, end, count); + }, + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(a) => if let (Some(lo), Some(hi)) = (low.try_u64(), high.try_u64()) { + coord_window!(a.data, lo, hi, u64, start, end, count); + }, + NumericArray::UInt32(a) => if let (Some(lo), Some(hi)) = (low.try_u64(), high.try_u64()) { + coord_window!(a.data, lo, hi, u64, start, end, count); + }, + NumericArray::UInt64(a) => if let (Some(lo), Some(hi)) = (low.try_u64(), high.try_u64()) { + coord_window!(a.data, lo, hi, u64, start, end, count); + }, + NumericArray::Float32(a) => if let (Some(lo), Some(hi)) = (low.try_f64(), high.try_f64()) { + coord_window!(a.data, lo, hi, f64, start, end, count); + }, + NumericArray::Float64(a) => if let (Some(lo), Some(hi)) = (low.try_f64(), high.try_f64()) { + coord_window!(a.data, lo, hi, f64, start, end, count); + }, + NumericArray::Null => {} + }, + #[cfg(feature = "datetime")] + Array::TemporalArray(temporal) => { + let lo = low.try_i64().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "datetime ticks", + message: Some(format!("axis '{}' holds datetime coordinates", self.name)), + })?; + let hi = high.try_i64().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "datetime ticks", + message: Some(format!("axis '{}' holds datetime coordinates", self.name)), + })?; + match temporal { + TemporalArray::Datetime32(a) => { coord_window!(a.data, lo, hi, i64, start, end, count); } + TemporalArray::Datetime64(a) => { coord_window!(a.data, lo, hi, i64, start, end, count); } + TemporalArray::Null => {} + } + } + _ => return Err(MinarrowError::TypeError { + from: "coordinate array", to: "comparable value", + message: Some(format!( + "axis '{}' coordinates do not support range lookup", self.name + )), + }), + } + if start >= end { + return Err(MinarrowError::IndexError(format!( + "no values in [{:?}, {:?}] on axis '{}'", low, high, self.name + ))); + } + // A span wider than its match count contains coordinates outside + // the bounds, so a window would return wrong rows. + if end - start != count { + return Err(MinarrowError::IndexError(format!( + "coordinates on axis '{}' are not monotonic over [{:?}, {:?}], sort the axis or gather by position instead", + self.name, low, high + ))); + } + Ok((start, end)) + } + + /// Position of the coordinate closest to `value`. Numeric coordinates + /// measure distance in their native domain (i64, u64, or f64), and + /// datetime coordinates as tick difference. String and categorical + /// coordinates have no distance metric and return an error - use + /// exact lookup instead. Ties resolve to the earliest position. + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + pub(crate) fn try_coord_nearest(&self, value: &Scalar) -> Result { + let coords = self.coords.as_ref().ok_or_else(|| MinarrowError::ShapeError { + message: format!("axis '{}' has no coordinates for nearest lookup", self.name), + })?; + let found = match coords { + Array::TextArray(_) => return Err(MinarrowError::TypeError { + from: "text", to: "numeric", + message: Some(format!( + "axis '{}' holds text coordinates with no distance metric - use at", + self.name + )), + }), + Array::NumericArray(num) => match num { + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int8(a) => + value.try_i64().and_then(|t| coord_nearest!(a.data, |v| (v as i64).abs_diff(t))), + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int16(a) => + value.try_i64().and_then(|t| coord_nearest!(a.data, |v| (v as i64).abs_diff(t))), + NumericArray::Int32(a) => + value.try_i64().and_then(|t| coord_nearest!(a.data, |v| (v as i64).abs_diff(t))), + NumericArray::Int64(a) => + value.try_i64().and_then(|t| coord_nearest!(a.data, |v| v.abs_diff(t))), + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt8(a) => + value.try_u64().and_then(|t| coord_nearest!(a.data, |v| (v as u64).abs_diff(t))), + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(a) => + value.try_u64().and_then(|t| coord_nearest!(a.data, |v| (v as u64).abs_diff(t))), + NumericArray::UInt32(a) => + value.try_u64().and_then(|t| coord_nearest!(a.data, |v| (v as u64).abs_diff(t))), + NumericArray::UInt64(a) => + value.try_u64().and_then(|t| coord_nearest!(a.data, |v| v.abs_diff(t))), + NumericArray::Float32(a) => + value.try_f64().and_then(|t| coord_nearest!(a.data, |v| (v as f64 - t).abs())), + NumericArray::Float64(a) => + value.try_f64().and_then(|t| coord_nearest!(a.data, |v| (v - t).abs())), + NumericArray::Null => None, + }, + #[cfg(feature = "datetime")] + Array::TemporalArray(temporal) => { + let target = value.try_i64().ok_or_else(|| MinarrowError::TypeError { + from: "scalar", to: "datetime ticks", + message: Some(format!("axis '{}' holds datetime coordinates", self.name)), + })?; + match temporal { + TemporalArray::Datetime32(a) => + coord_nearest!(a.data, |t| (t as i64).abs_diff(target)), + TemporalArray::Datetime64(a) => + coord_nearest!(a.data, |t| t.abs_diff(target)), + TemporalArray::Null => None, + } + } + _ => return Err(MinarrowError::TypeError { + from: "coordinate array", to: "comparable value", + message: Some(format!( + "axis '{}' coordinates do not support nearest lookup", self.name + )), + }), + }; + found.ok_or_else(|| MinarrowError::IndexError(format!( + "axis '{}' has no comparable coordinates", self.name + ))) + } +} + +// **************************************************************** +// Trait implementations +// **************************************************************** + +/// Positional selection across every axis at once, delegating to `slice`. +/// The result is an unlabelled view. For named-axis selection with the +/// labels carried through, use `.select()`. +#[cfg(all(feature = "views", feature = "select"))] +impl AxisSelection for XArray { + type View = NdArrayV; + + fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV { + self.slice(selection) + } + + fn get_axis_count(&self) -> usize { + self.ndim() + } +} + +impl Shape for XArray { + fn shape(&self) -> ShapeDim { + match &self.data { + NdArrayE::Owned(nd) => Shape::shape(nd), + #[cfg(feature = "views")] + NdArrayE::View(v) => Shape::shape(v), + } + } +} + +impl Concatenate for XArray { + fn concat(self, other: Self) -> Result { + if self.axes.len() != other.axes.len() { + return Err(MinarrowError::IncompatibleTypeError { + from: "XArray", to: "XArray", + message: Some(format!("{} dims vs {} dims", self.axes.len(), other.axes.len())), + }); + } + for (a, b) in self.axes.iter().zip(other.axes.iter()) { + if a.name != b.name { + return Err(MinarrowError::IncompatibleTypeError { + from: "XArray", to: "XArray", + message: Some(format!("dim name mismatch: '{}' vs '{}'", a.name, b.name)), + }); + } + } + + // Destructure to avoid cloning axes + let XArray { data: self_data, axes: self_axes } = self; + let XArray { data: other_data, axes: other_axes } = other; + + let to_ndarray = |data: NdArrayE| -> NdArray { + match data { + NdArrayE::Owned(nd) => nd, + #[cfg(feature = "views")] + NdArrayE::View(v) => v.to_ndarray(), + } + }; + let self_nd = to_ndarray(self_data); + let other_nd = to_ndarray(other_data); + let new_data = self_nd.concat(other_nd)?; + + // Axis 0 coords concatenate when both sides carry them. Non-0 + // axes describe the same positions on both sides, so their + // coords must agree, and a side missing them adopts the other's. + let mut new_axes = Vec::with_capacity(self_axes.len()); + for (i, (a, b)) in self_axes.into_iter().zip(other_axes.into_iter()).enumerate() { + if i == 0 { + let merged_coords = match (a.coords, b.coords) { + (Some(ca), Some(cb)) => Some(ca.concat(cb)?), + (None, None) => None, + _ => { + return Err(MinarrowError::IncompatibleTypeError { + from: "XArray", + to: "XArray", + message: Some(format!( + "axis '{}' is labelled on one side only, assign_coords or drop_coords first", + a.name + )), + }); + } + }; + new_axes.push(Axis { name: a.name, coords: merged_coords }); + } else { + let coords = match (a.coords, b.coords) { + (Some(ca), Some(cb)) => { + if ca != cb { + return Err(MinarrowError::IncompatibleTypeError { + from: "XArray", + to: "XArray", + message: Some(format!( + "axis '{}' coordinates differ between the two sides", + a.name + )), + }); + } + Some(ca) + } + (Some(ca), None) => Some(ca), + (None, Some(cb)) => Some(cb), + (None, None) => None, + }; + new_axes.push(Axis { name: a.name, coords }); + } + } + + Ok(XArray { data: NdArrayE::Owned(new_data), axes: new_axes }) + } +} + +impl<'a, T: Float> IntoIterator for &'a XArray { + type Item = T; + type IntoIter = Box + 'a>; + + fn into_iter(self) -> Self::IntoIter { + match &self.data { + NdArrayE::Owned(nd) => Box::new(nd.into_iter()), + #[cfg(feature = "views")] + NdArrayE::View(v) => Box::new(v.into_iter()), + } + } +} + +impl PartialEq for XArray { + fn eq(&self, other: &Self) -> bool { + if self.axes != other.axes { return false; } + if self.shape() != other.shape() { return false; } + self.into_iter().zip(other.into_iter()).all(|(a, b)| a == b) + } +} + +impl fmt::Debug for XArray { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let shape = self.shape(); + let dims: Vec = self.axes.iter() + .zip(shape.iter()) + .map(|(ax, &size)| { + let label = if ax.coords.is_some() { " (labelled)" } else { "" }; + format!("{}={}{}", ax.name, size, label) + }) + .collect(); + let storage = match &self.data { + NdArrayE::Owned(_) => "owned", + #[cfg(feature = "views")] + NdArrayE::View(_) => "view", + }; + write!(f, "XArray [{}] ({})", dims.join(", "), storage) + } +} + + +// **************************************************************** +// TryFrom
+// **************************************************************** + +impl TryFrom
for XArray { + type Error = MinarrowError; + + fn try_from(table: Table) -> Result { + let col_names: Vec = table.col_names().iter().map(|s| s.to_string()).collect(); + let table_name = table.name.clone(); + let data = NdArray::try_from(&table)?; + + let obs_axis = Axis::named( + if table_name.is_empty() { "observation" } else { &table_name } + ); + let feat_axis = Axis::with_coords( + "feature", + Array::from_string32(StringArray::from_slice( + &col_names.iter().map(|s| s.as_str()).collect::>() + )), + ); + + Ok(XArray { data: NdArrayE::Owned(data), axes: vec![obs_axis, feat_axis] }) + } +} + +// **************************************************************** +// Tests +// **************************************************************** + +#[cfg(test)] +mod tests { + use super::*; + #[cfg(all(feature = "views", feature = "select"))] + use crate::nd; + use std::sync::Arc; + use crate::{FloatArray, NumericArray, FieldArray}; + + #[cfg(all(feature = "views", feature = "select"))] + #[test] + fn axis_selection_positional() { + let xa = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]), + &["obs", "feat"], + ); + let v = xa.s(nd![0..2, 1]); + assert_eq!(v.shape(), &[2]); + assert_eq!(v.get(&[0]), 4.0); + assert_eq!(v.get(&[1]), 5.0); + } + + #[test] + fn apply_preserves_axes() { + let xa = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + let out = xa.apply(|x| x * 2.0); + assert_eq!(out.dim_names(), vec!["obs", "feat"]); + assert_eq!(out.get(&[1, 1]), 8.0); + assert_eq!(xa.get(&[1, 1]), 4.0); + } + + #[test] + fn apply_mut_owned() { + let mut xa = XArray::new(NdArray::from_slice(&[1.0, 2.0], &[2]), &["t"]); + xa.apply_mut(|x| x + 10.0); + assert_eq!(xa.get(&[0]), 11.0); + assert_eq!(xa.get(&[1]), 12.0); + } + + #[cfg(feature = "views")] + #[test] + #[should_panic(expected = "cannot mutate a view")] + fn apply_mut_view_panics() { + let base = XArray::new(NdArray::from_slice(&[1.0, 2.0], &[2]), &["t"]); + let mut view = base.select(&[("t", &(0..2))]); + view.apply_mut(|x| x + 1.0); + } + + fn make_2d() -> NdArray { + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]) + } + + #[test] + fn f32_element_type() { + let data = NdArray::::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]); + let xa = XArray::new(data, &["obs", "feat"]); + assert_eq!(xa.get(&[0, 0]), 1.0f32); + assert_eq!(xa.get(&[2, 1]), 6.0f32); + let sub = xa.select(&[("obs", &(0..2))]); + assert_eq!(sub.shape(), vec![2, 2]); + assert_eq!(sub.get(&[1, 1]), 5.0f32); + } + + fn float_coords(vals: &[f64]) -> Array { + Array::NumericArray(NumericArray::Float64( + Arc::new(FloatArray::from_slice(vals)) + )) + } + + // *** Construction ************************************************ + + #[test] + fn rank_zero_scalar() { + let xa = XArray::new(NdArray::from_slice(&[5.0], &[]), &[]); + assert_eq!(xa.ndim(), 0); + assert_eq!(xa.shape(), Vec::::new()); + assert!(xa.dim_names().is_empty()); + assert_eq!(xa.len(), 1); + assert_eq!(xa.get(&[]), 5.0); + } + + #[test] + fn new_with_dim_names() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + assert_eq!(xa.ndim(), 2); + assert_eq!(xa.dim_names(), vec!["obs", "feat"]); + } + + #[test] + fn from_ndarray_auto_names() { + let xa = XArray::from_ndarray(make_2d()); + assert_eq!(xa.dim_names(), vec!["dim_0", "dim_1"]); + } + + #[test] + fn with_axes_and_coords() { + let axes = vec![ + Axis::named("obs"), + Axis::with_coords("feat", float_coords(&[10.0, 20.0])), + ]; + let xa = XArray::with_axes(make_2d(), axes); + assert!(xa.ax("feat").coords.is_some()); + assert!(xa.ax("obs").coords.is_none()); + } + + // *** Axis access ************************************************* + + #[test] + fn ax_and_dim() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + assert_eq!(xa.ax("feat").name, "feat"); + assert_eq!(xa.dim("obs"), 0); + assert_eq!(xa.dim("feat"), 1); + } + + #[test] + fn try_ax_and_try_dim() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + assert!(xa.try_ax("feat").is_some()); + assert!(xa.try_ax("missing").is_none()); + assert_eq!(xa.try_dim("obs"), Some(0)); + assert_eq!(xa.try_dim("missing"), None); + } + + // *** Delegation ************************************************** + + #[test] + fn delegates_shape_and_access() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + assert_eq!(xa.shape(), vec![3, 2]); + assert_eq!(xa.len(), 6); + assert_eq!(xa.get(&[0, 0]), 1.0); + let obs0: Vec = (&xa.obs(0)).into_iter().collect(); + assert_eq!(obs0, vec![1.0, 4.0]); + } + + #[test] + fn iteration() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let vals: Vec = (&xa).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + } + + // *** Label management ******************************************** + + #[test] + fn rename_dim() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.rename_dim("feat", "variable"); + assert_eq!(xa.dim("variable"), 1); + } + + #[test] + fn assign_and_drop_coords() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + assert!(xa.ax("feat").coords.is_none()); + xa.assign_coords("feat", float_coords(&[10.0, 20.0])); + assert!(xa.ax("feat").coords.is_some()); + xa.drop_coords("feat"); + assert!(xa.ax("feat").coords.is_none()); + } + + // *** Positional selection **************************************** + + #[cfg(feature = "views")] + #[test] + fn select_range() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let result = xa.select(&[("obs", &(0..2))]); + assert_eq!(result.shape(), &[2, 2]); + assert_eq!(result.dim_names(), vec!["obs", "feat"]); + assert!(!result.is_owned()); + } + + #[cfg(feature = "views")] + #[test] + fn select_index_collapses() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let result = xa.select(&[("feat", &0)]); + assert_eq!(result.ndim(), 1); + assert_eq!(result.dim_names(), vec!["obs"]); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0]); + } + + #[cfg(feature = "views")] + #[test] + fn select_multi_axis() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let result = xa.select(&[("obs", &(1..3)), ("feat", &1)]); + assert_eq!(result.ndim(), 1); + assert_eq!(result.dim_names(), vec!["obs"]); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![5.0, 6.0]); + } + + #[cfg(feature = "views")] + #[test] + fn select_preserves_coords() { + let mut xa = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[4, 2]), + &["obs", "feat"], + ); + xa.assign_coords("obs", float_coords(&[10.0, 20.0, 30.0, 40.0])); + let result = xa.select(&[("obs", &(1..3))]); + let coords = result.ax("obs").coords.as_ref().unwrap(); + assert_eq!(coords.len(), 2); + } + + // *** Coordinate selection **************************************** + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn at_by_coord() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("feat", float_coords(&[10.0, 20.0])); + let result = xa.at("feat", 20.0); + assert_eq!(result.ndim(), 1); + assert_eq!(result.dim_names(), vec!["obs"]); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![4.0, 5.0, 6.0]); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn try_at_not_found() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("feat", float_coords(&[10.0, 20.0])); + assert!(xa.try_at("feat", 99.0).is_err()); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn between_coord_range() { + let data = NdArray::from_slice( + &[1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0, 30.0, 40.0, 50.0], + &[5, 2], + ); + let mut xa = XArray::new(data, &["obs", "feat"]); + xa.assign_coords("obs", float_coords(&[0.0, 1.0, 2.0, 3.0, 4.0])); + let result = xa.between("obs", 1.0, 3.0); + assert_eq!(result.shape(), vec![3, 2]); + let vals: Vec = (&result).into_iter().take(3).collect(); + assert_eq!(vals, vec![2.0, 3.0, 4.0]); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn try_between_empty_range() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("obs", float_coords(&[0.0, 1.0, 2.0])); + assert!(xa.try_between("obs", 100.0, 200.0).is_err()); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn at_string_coords() { + use crate::arr_str32; + let mut xa = XArray::new(make_2d(), &["ticker", "feat"]); + xa.assign_coords("ticker", arr_str32!(&["ES", "NQ", "CL"])); + let result = xa.at("ticker", "NQ"); + assert_eq!(result.ndim(), 1); + assert_eq!(result.dim_names(), vec!["feat"]); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + assert!(xa.try_at("ticker", "ZB").is_err()); + } + + #[cfg(all( + feature = "views", + feature = "select", + feature = "scalar_type", + any(not(feature = "default_categorical_8"), feature = "extended_categorical") + ))] + #[test] + fn at_categorical_coords() { + use crate::arr_cat32; + let mut xa = XArray::new(make_2d(), &["ticker", "feat"]); + xa.assign_coords("ticker", arr_cat32!(&["ES", "NQ", "ES"])); + let result = xa.at("ticker", "NQ"); + assert_eq!(result.ndim(), 1); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + assert!(xa.try_at("ticker", "ZB").is_err()); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn coord_lookup_integer_coords() { + use crate::arr_i64; + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("obs", arr_i64![100, 200, 400]); + let result = xa.at("obs", 200); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + let ranged = xa.between("obs", 150, 450); + assert_eq!(ranged.shape(), vec![2, 2]); + let near = xa.nearest("obs", 230); + let vals: Vec = (&near).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn between_string_coords() { + use crate::arr_str32; + let mut xa = XArray::new(make_2d(), &["ticker", "feat"]); + xa.assign_coords("ticker", arr_str32!(&["alpha", "bravo", "charlie"])); + // Lexicographic inclusive bounds. + let result = xa.between("ticker", "bravo", "delta"); + assert_eq!(result.shape(), vec![2, 2]); + let vals: Vec = (&result).into_iter().take(2).collect(); + assert_eq!(vals, vec![2.0, 3.0]); + } + + #[cfg(all( + feature = "views", + feature = "select", + feature = "scalar_type", + feature = "datetime" + ))] + #[test] + fn at_datetime_coords() { + use crate::arr_dt64; + use crate::enums::time_units::TimeUnit; + let mut xa = XArray::new(make_2d(), &["time", "feat"]); + xa.assign_coords("time", arr_dt64!(TimeUnit::Milliseconds; 1_000, 2_000, 3_000)); + let result = xa.at("time", 2_000); + assert_eq!(result.ndim(), 1); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + assert!(xa.try_at("time", 5_000).is_err()); + } + + #[cfg(all( + feature = "views", + feature = "select", + feature = "scalar_type", + feature = "datetime" + ))] + #[test] + fn between_datetime_coords() { + use crate::arr_dt64; + use crate::enums::time_units::TimeUnit; + let mut xa = XArray::new(make_2d(), &["time", "feat"]); + xa.assign_coords("time", arr_dt64!(TimeUnit::Milliseconds; 1_000, 2_000, 3_000)); + let result = xa.between("time", 1_500, 3_500); + assert_eq!(result.shape(), vec![2, 2]); + let vals: Vec = (&result).into_iter().take(2).collect(); + assert_eq!(vals, vec![2.0, 3.0]); + } + + #[cfg(all( + feature = "views", + feature = "select", + feature = "scalar_type", + feature = "datetime" + ))] + #[test] + fn nearest_datetime_coords() { + use crate::arr_dt64; + use crate::enums::time_units::TimeUnit; + let mut xa = XArray::new(make_2d(), &["time", "feat"]); + xa.assign_coords("time", arr_dt64!(TimeUnit::Milliseconds; 1_000, 2_000, 4_000)); + let result = xa.nearest("time", 2_700); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn nearest_numeric_coords() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("obs", float_coords(&[10.0, 20.0, 40.0])); + let result = xa.nearest("obs", 24.0); + assert_eq!(result.ndim(), 1); + let vals: Vec = (&result).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + // Ties resolve to the earliest position. + let tied = xa.nearest("obs", 30.0); + let vals: Vec = (&tied).into_iter().collect(); + assert_eq!(vals, vec![2.0, 5.0]); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn nearest_text_coords_rejected() { + use crate::arr_str32; + let mut xa = XArray::new(make_2d(), &["ticker", "feat"]); + xa.assign_coords("ticker", arr_str32!(&["a", "b", "c"])); + assert!(xa.try_nearest("ticker", "b").is_err()); + } + + #[test] + fn axis_equality_includes_coords() { + let a = Axis::with_coords("obs", float_coords(&[1.0, 2.0])); + let b = Axis::with_coords("obs", float_coords(&[1.0, 2.0])); + let c = Axis::with_coords("obs", float_coords(&[1.0, 3.0])); + assert_eq!(a, b); + assert_ne!(a, c); + assert_ne!(a, Axis::named("obs")); + } + + // *** Transpose *************************************************** + + #[test] + fn transpose_reorders_axes() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let t = xa.transpose(&["feat", "obs"]).unwrap(); + assert_eq!(t.dim_names(), vec!["feat", "obs"]); + assert_eq!(t.shape(), &[2, 3]); + } + + // *** Concatenate ************************************************* + + #[test] + fn concat_matching_dims() { + let a = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + let b = XArray::new( + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + &["obs", "feat"], + ); + let c = a.concat(b).unwrap(); + assert_eq!(c.shape(), &[4, 2]); + assert_eq!(c.dim_names(), vec!["obs", "feat"]); + } + + #[test] + fn concat_merges_axis0_coords() { + let mut a = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + a.assign_coords("obs", float_coords(&[0.0, 1.0])); + let mut b = XArray::new( + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + &["obs", "feat"], + ); + b.assign_coords("obs", float_coords(&[2.0, 3.0])); + let c = a.concat(b).unwrap(); + let coords = c.ax("obs").coords.as_ref().unwrap(); + assert_eq!(coords.len(), 4); + } + + #[test] + fn concat_mismatched_dims_fails() { + let a = XArray::new(make_2d(), &["obs", "feat"]); + let b = XArray::new(make_2d(), &["obs", "variable"]); + assert!(a.concat(b).is_err()); + } + + // *** Table roundtrip ********************************************* + + #[test] + fn try_from_table() { + let c0 = FieldArray::from_arr("height", Array::NumericArray( + NumericArray::Float64(Arc::new(FloatArray::from_slice(&[175.0, 168.0]))) + )); + let c1 = FieldArray::from_arr("weight", Array::NumericArray( + NumericArray::Float64(Arc::new(FloatArray::from_slice(&[82.0, 71.0]))) + )); + let table = Table::new("patients".to_string(), Some(vec![c0, c1])); + let xa = XArray::try_from(table).unwrap(); + assert_eq!(xa.dim_names(), vec!["patients", "feature"]); + assert!(xa.ax("feature").coords.is_some()); + } + + #[test] + fn to_table_with_coord_names() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("feat", Array::from_string32( + StringArray::from_slice(&["height", "weight"]) + )); + let table = xa.to_table().unwrap(); + assert_eq!(table.col_names(), vec!["height", "weight"]); + } + + // *** Misc ******************************************************** + + #[test] + fn clone_and_eq() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + assert_eq!(xa, xa.clone()); + } + + #[test] + fn debug_format() { + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("feat", float_coords(&[10.0, 20.0])); + let s = format!("{:?}", xa); + assert!(s.contains("obs=3")); + assert!(s.contains("feat=2 (labelled)")); + assert!(s.contains("owned")); + } + + #[cfg(feature = "views")] + #[test] + fn select_produces_view_debug() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let result = xa.select(&[("obs", &(0..2))]); + assert!(!result.is_owned()); + assert!(format!("{:?}", result).contains("view")); + } + + #[test] + fn to_owned_from_view() { + #[cfg(feature = "views")] + { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let view = xa.select(&[("obs", &(0..2))]); + assert!(!view.is_owned()); + let owned = view.to_owned(); + assert!(owned.is_owned()); + assert_eq!(owned.shape(), &[2, 2]); + } + } + + #[test] + fn three_dimensional() { + let data = NdArray::from_slice( + &(1..=24).map(|x| x as f64).collect::>(), + &[2, 3, 4], + ); + let xa = XArray::new(data, &["obs", "feat", "time"]); + assert_eq!(xa.ndim(), 3); + assert_eq!(xa.dim("time"), 2); + } + + #[test] + fn transpose_current_order_is_identity() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let same = xa.transpose(&["obs", "feat"]).unwrap(); + assert_eq!(same.dim_names(), vec!["obs", "feat"]); + assert_eq!(same.shape(), vec![3, 2]); + assert_eq!(same.get(&[2, 1]), xa.get(&[2, 1])); + } + + #[test] + fn transpose_reversed_moves_data() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let t = xa.transpose(&["feat", "obs"]).unwrap(); + assert_eq!(t.dim_names(), vec!["feat", "obs"]); + assert_eq!(t.get(&[0, 2]), xa.get(&[2, 0])); + assert_eq!(t.get(&[1, 0]), xa.get(&[0, 1])); + } + + #[test] + fn transpose_duplicate_dim_errors() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let err = xa.transpose(&["obs", "obs"]).unwrap_err(); + assert!(err.to_string().contains("must be distinct")); + } + + #[test] + fn transpose_unknown_dim_errors() { + let xa = XArray::new(make_2d(), &["obs", "feat"]); + let err = xa.transpose(&["obs", "missing"]).unwrap_err(); + assert!(err.to_string().contains("no axis named 'missing'")); + } + + #[cfg(all(feature = "views", feature = "select", feature = "scalar_type"))] + #[test] + fn try_between_unsorted_coords_errors() { + // The covering span 0..3 includes the out-of-bounds 100.0, so a + // window would return wrong rows. + let mut xa = XArray::new(make_2d(), &["obs", "feat"]); + xa.assign_coords("obs", float_coords(&[20.0, 100.0, 21.0])); + let err = xa.try_between("obs", 15.0, 25.0).unwrap_err(); + assert!(err.to_string().contains("not monotonic")); + + // Sorted coordinates over the same bounds window cleanly. + let mut sorted = XArray::new(make_2d(), &["obs", "feat"]); + sorted.assign_coords("obs", float_coords(&[20.0, 21.0, 100.0])); + let result = sorted.try_between("obs", 15.0, 25.0).unwrap(); + assert_eq!(result.shape(), vec![2, 2]); + } + + #[test] + fn concat_axis0_coords_on_one_side_fails() { + let mut a = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + a.assign_coords("obs", float_coords(&[0.0, 1.0])); + let b = XArray::new( + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + &["obs", "feat"], + ); + let err = a.concat(b).unwrap_err(); + assert!(err.to_string().contains("labelled on one side")); + } + + #[test] + fn concat_non_zero_axis_coord_mismatch_fails() { + let mut a = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + a.assign_coords("feat", float_coords(&[10.0, 20.0])); + let mut b = XArray::new( + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + &["obs", "feat"], + ); + b.assign_coords("feat", float_coords(&[10.0, 30.0])); + let err = a.concat(b).unwrap_err(); + assert!(err.to_string().contains("coordinates differ")); + } + + #[test] + fn concat_non_zero_axis_adopts_coords() { + let mut a = XArray::new( + NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]), + &["obs", "feat"], + ); + a.assign_coords("feat", float_coords(&[10.0, 20.0])); + let b = XArray::new( + NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]), + &["obs", "feat"], + ); + let c = a.concat(b).unwrap(); + let coords = c.ax("feat").coords.as_ref().unwrap(); + assert_eq!(*coords, float_coords(&[10.0, 20.0])); + } + + #[cfg(feature = "views")] + #[test] + #[should_panic(expected = "coords but dimension size is")] + fn from_view_coord_length_mismatch_panics() { + let nd = make_2d(); + let axes = vec![ + Axis::named("obs"), + Axis::with_coords("feat", float_coords(&[10.0])), + ]; + let _ = XArray::from_view(nd.as_view(), axes); + } + +} diff --git a/src/traits/byte_size.rs b/src/traits/byte_size.rs index b7e18cd..433e1f3 100644 --- a/src/traits/byte_size.rs +++ b/src/traits/byte_size.rs @@ -328,6 +328,82 @@ impl ByteSize for Matrix { } } +/// ByteSize for NdArray (when ndarray feature is enabled) +#[cfg(feature = "ndarray")] +use crate::structs::ndarray::NdArray; + +#[cfg(feature = "ndarray")] +impl ByteSize for NdArray { + fn est_bytes(&self) -> usize { + // Physical backing buffer, including any stride padding. + self.data.est_bytes() + } +} + +/// ByteSize for NdArrayV - proportional estimate from the backing array +#[cfg(all(feature = "ndarray", feature = "views"))] +use crate::NdArrayV; + +#[cfg(all(feature = "ndarray", feature = "views"))] +impl ByteSize for NdArrayV { + fn est_bytes(&self) -> usize { + let full_len = self.source.len(); + let full_bytes = self.source.est_bytes(); + if full_len > 0 { + (full_bytes * self.len()) / full_len + } else { + 0 + } + } +} + +/// ByteSize for SuperNdArray - sum of batch estimates +#[cfg(all(feature = "ndarray", feature = "chunked"))] +use crate::SuperNdArray; + +#[cfg(all(feature = "ndarray", feature = "chunked"))] +impl ByteSize for SuperNdArray { + fn est_bytes(&self) -> usize { + self.batches.iter().map(|batch| batch.est_bytes()).sum() + } +} + +/// ByteSize for SuperNdArrayV - sum of slice estimates +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +use crate::SuperNdArrayV; + +#[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] +impl ByteSize for SuperNdArrayV { + fn est_bytes(&self) -> usize { + self.slices.iter().map(|slice| slice.est_bytes()).sum() + } +} + +/// ByteSize for XArray - storage plus coordinate arrays +#[cfg(feature = "xarray")] +use crate::XArray; + +#[cfg(feature = "xarray")] +impl ByteSize for XArray { + fn est_bytes(&self) -> usize { + use crate::structs::xarray::NdArrayE; + let data_bytes = match self.storage() { + NdArrayE::Owned(nd) => nd.est_bytes(), + #[cfg(feature = "views")] + NdArrayE::View(v) => v.est_bytes(), + }; + let coord_bytes: usize = self + .axes() + .iter() + .map(|axis| { + axis.name.capacity() + + axis.coords.as_ref().map(|c| c.est_bytes()).unwrap_or(0) + }) + .sum(); + data_bytes + coord_bytes + } +} + /// ByteSize for Cube (when cube feature is enabled) #[cfg(feature = "cube")] use crate::Cube; @@ -414,6 +490,16 @@ impl ByteSize for Value { Value::FieldArray(fa) => fa.est_bytes(), #[cfg(feature = "matrix")] Value::Matrix(m) => m.est_bytes(), + #[cfg(feature = "ndarray")] + Value::NdArray(nd) => nd.est_bytes(), + #[cfg(all(feature = "ndarray", feature = "views"))] + Value::NdArrayView(v) => v.est_bytes(), + #[cfg(all(feature = "ndarray", feature = "chunked"))] + Value::SuperNdArray(snd) => snd.est_bytes(), + #[cfg(all(feature = "ndarray", feature = "chunked", feature = "views"))] + Value::SuperNdArrayView(sv) => sv.est_bytes(), + #[cfg(feature = "xarray")] + Value::XArray(xa) => xa.est_bytes(), #[cfg(feature = "cube")] Value::Cube(c) => c.est_bytes(), Value::VecValue(vec) => { diff --git a/src/traits/selection.rs b/src/traits/selection.rs index 077d5ea..7d5914c 100644 --- a/src/traits/selection.rs +++ b/src/traits/selection.rs @@ -56,6 +56,27 @@ pub trait DataSelector { fn is_contiguous(&self) -> bool { false // Default: assume non-contiguous } + + /// Resolve this selector against one axis of `dim_size`, producing a + /// `(start, end, collapse)` span. Ranges window the axis and a single + /// index collapses it. Panics when the selection falls outside the + /// axis. Index arrays have no strided-view representation, so they + /// are rejected - gather through [`RowSelection::r`] instead. + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + // usize::MAX so out-of-range indices survive to the bounds check + // below rather than being silently filtered. + let indices = self.resolve_indices(usize::MAX); + assert_eq!( + indices.len(), 1, + "axis selection: index arrays take a single index or a contiguous range; gather with r() instead" + ); + assert!( + indices[0] < dim_size, + "axis selection: index {} out of bounds (size {})", indices[0], dim_size + ); + (indices[0], indices[0] + 1, true) + } } // These traits are implemented on structures like Table, ArrayV, etc. @@ -78,7 +99,10 @@ pub trait ColumnSelection { /// Select fields/columns by name, index, or range /// + /// Shorthand alias for `col` + /// /// # Examples + /// ```ignore /// table.c("age") // single column by name /// table.c(&["a", "b"]) // multiple columns by name /// table.c(0) // single column by index @@ -86,7 +110,18 @@ pub trait ColumnSelection { /// ``` fn c(&self, selection: S) -> Self::View; - /// Alias for `c` - select column by name + /// Select a single column by name + /// + /// Named form of `c` - use `c` for selection by index, range, + /// or multiple names. + /// + /// # Examples + /// ```ignore + /// table.col("age") // single column by name + /// table.c(&["a", "b"]) // multiple columns by name + /// table.c(0) // single column by index + /// table.c(0..3) // columns by range + /// ``` fn col(&self, name: &str) -> Self::View { self.c(name) } @@ -111,17 +146,88 @@ pub trait RowSelection { /// Select rows by index or range /// + /// Shorthand alias for `row` + /// /// # Examples + /// ```ignore /// table.r(5) // single row /// table.r(&[1, 3, 5]) // specific rows /// table.r(0..10) // row range /// ``` fn r(&self, selection: S) -> Self::View; + /// Select a single row by index + /// + /// Named form of `r` - use `r` for selection by range or + /// index array. + /// + /// # Examples + /// ```ignore + /// table.row(5) // single row + /// table.r(&[1, 3, 5]) // specific rows + /// table.r(0..10) // row range + /// ``` + fn row(&self, idx: usize) -> Self::View { + self.r(idx) + } + /// Get the count for data resolution fn get_row_count(&self) -> usize; } +/// Trait for types that support selection across any axis combination. +/// +/// The N-dimensional member of the selection family, alongside +/// [`ColumnSelection`] and [`RowSelection`]. Each axis takes any +/// [`DataSelector`] - a single index collapses the dimension, and a +/// contiguous range keeps it. +/// +/// # Examples +/// ```ignore +/// arr.s(&[&(1..4), &2]) // rows 1..4 of column 2 (2D) +/// arr.s(nd![1..4, 2]) // the same through the nd! macro +/// arr.s(nd![0..2, 1, 0..3]) // mixed selection (3D) +/// ``` +#[cfg(feature = "ndarray")] +pub trait AxisSelection { + /// The view type returned by axis selection e.g. NdArrayV + type View; + + /// Select along every axis at once, one [`DataSelector`] per axis + /// + /// Shorthand alias for `select` + /// + /// A single index collapses its dimension, and a contiguous range + /// keeps it. + /// + /// # Examples + /// ```ignore + /// arr.s(nd![1..4, 2]) // rows 1..4 of column 2 (2D) + /// arr.s(nd![0..2, 1, 0..3]) // mixed selection (3D) + /// arr.s(nd![.., 5]) // full range on axis 0 + /// arr.s(&[&(1..4), &2]) // without the nd! macro + /// ``` + fn s(&self, selection: &[&dyn DataSelector]) -> Self::View; + + /// Select along every axis at once, one [`DataSelector`] per axis + /// + /// A single index collapses its dimension, and a contiguous range + /// keeps it. + /// + /// # Examples + /// ```ignore + /// arr.select(nd![1..4, 2]) // rows 1..4 of column 2 (2D) + /// arr.select(nd![0..2, 1, 0..3]) // mixed selection (3D) + /// arr.select(nd![.., 5]) // full range on axis 0 + /// ``` + fn select(&self, selection: &[&dyn DataSelector]) -> Self::View { + self.s(selection) + } + + /// Get the axis count for selection resolution + fn get_axis_count(&self) -> usize; +} + /// Combined trait for 2D selection (field + data dimensions) /// /// This trait is automatically implemented for any type that implements @@ -328,6 +434,46 @@ impl DataSelector for usize { Vec::new() } } + + /// A single index is a contiguous window of length one. + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + *self < dim_size, + "axis selection: index {} out of bounds (size {})", self, dim_size + ); + (*self, *self + 1, true) + } +} + +/// Single data index from a plain integer literal. Negative values +/// resolve to nothing. +impl DataSelector for i32 { + fn resolve_indices(&self, count: usize) -> Vec { + if *self >= 0 && (*self as usize) < count { + vec![*self as usize] + } else { + Vec::new() + } + } + + /// A single index is a contiguous window of length one. + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + *self >= 0 && (*self as usize) < dim_size, + "axis selection: index {} out of bounds (size {})", self, dim_size + ); + (*self as usize, *self as usize + 1, true) + } } /// Multiple data indices @@ -361,6 +507,40 @@ impl DataSelector for Range { fn is_contiguous(&self) -> bool { true } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.start <= self.end && self.end <= dim_size, + "axis selection: range {}..{} out of bounds (size {})", + self.start, self.end, dim_size + ); + (self.start, self.end, false) + } +} + +/// Data range selection from plain integer literals. Negative bounds +/// clamp to zero. +impl DataSelector for Range { + fn resolve_indices(&self, count: usize) -> Vec { + let start = self.start.max(0) as usize; + let end = (self.end.max(0) as usize).min(count); + (start..end).collect() + } + + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.start >= 0 && self.end >= self.start && (self.end as usize) <= dim_size, + "axis selection: range {}..{} out of bounds (size {})", + self.start, self.end, dim_size + ); + (self.start as usize, self.end as usize, false) + } } /// Data range from selection @@ -372,6 +552,36 @@ impl DataSelector for RangeFrom { fn is_contiguous(&self) -> bool { true } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.start <= dim_size, + "axis selection: range {}.. out of bounds (size {})", self.start, dim_size + ); + (self.start, dim_size, false) + } +} + +/// Data range from selection from a plain integer literal. +impl DataSelector for RangeFrom { + fn resolve_indices(&self, count: usize) -> Vec { + let start = self.start.max(0) as usize; + (start..count).collect() + } + + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.start >= 0 && (self.start as usize) <= dim_size, + "axis selection: range {}.. out of bounds (size {})", self.start, dim_size + ); + (self.start as usize, dim_size, false) + } } /// Data range to selection @@ -384,6 +594,37 @@ impl DataSelector for RangeTo { fn is_contiguous(&self) -> bool { true } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.end <= dim_size, + "axis selection: range ..{} out of bounds (size {})", self.end, dim_size + ); + (0, self.end, false) + } +} + +/// Data range to selection from a plain integer literal. Negative bounds +/// clamp to zero. +impl DataSelector for RangeTo { + fn resolve_indices(&self, count: usize) -> Vec { + let end = (self.end.max(0) as usize).min(count); + (0..end).collect() + } + + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.end >= 0 && (self.end as usize) <= dim_size, + "axis selection: range ..{} out of bounds (size {})", self.end, dim_size + ); + (0, self.end as usize, false) + } } /// Data full range selection @@ -395,6 +636,11 @@ impl DataSelector for RangeFull { fn is_contiguous(&self) -> bool { true } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + (0, dim_size, false) + } } /// Data inclusive range selection @@ -408,4 +654,41 @@ impl DataSelector for RangeInclusive { fn is_contiguous(&self) -> bool { true } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + self.start() <= self.end() && *self.end() < dim_size, + "axis selection: range {}..={} out of bounds (size {})", + self.start(), self.end(), dim_size + ); + (*self.start(), *self.end() + 1, false) + } +} + +/// Data inclusive range selection from plain integer literals. Negative +/// bounds clamp to zero. +impl DataSelector for RangeInclusive { + fn resolve_indices(&self, count: usize) -> Vec { + if *self.end() < 0 { + return Vec::new(); + } + let start = (*self.start()).max(0) as usize; + let end = (*self.end() as usize + 1).min(count); + (start..end).collect() + } + + fn is_contiguous(&self) -> bool { + true + } + + #[cfg(feature = "ndarray")] + fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) { + assert!( + *self.start() >= 0 && self.start() <= self.end() && (*self.end() as usize) < dim_size, + "axis selection: range {}..={} out of bounds (size {})", + self.start(), self.end(), dim_size + ); + (*self.start() as usize, *self.end() as usize + 1, false) + } }