Skip to main content

minarrow/structs/
ndarray.rs

1//! # `NdArray`
2//!
3//! An N-dimensional container for contiguous `f32` or `f64` data.
4//!
5//! `NdArray` is Minarrow's general-purpose interchange container for numeric
6//! data received from files, sensors, networks, Python, DLPack, or other
7//! runtimes. It supports construction, indexing, views, selection, iteration,
8//! conversion, and transfer to external numerical systems.
9//!
10//! It is not intended to provide a complete numerical-computing runtime.
11//! Computational workloads should be delegated to the application's preferred
12//! kernels, BLAS implementation, or machine-learning framework.
13//!
14//! ## Shape and element type
15//!
16//! `NdArray` is generic over `T: Float`, supporting `f32` and `f64`. Rank is
17//! determined at runtime. Shapes with one to five dimensions are stored inline,
18//! avoiding a heap allocation for the common case.
19//!
20//! Shape semantics follow standard tensor conventions:
21//!
22//! - `&[]` is a rank-zero array containing one logical value.
23//! - `[0]` is a rank-one array containing no values.
24//!
25//! ## Memory layout
26//!
27//! Data is stored in compact column-major order without padding between
28//! dimensions. The complete logical array therefore occupies one contiguous
29//! buffer, and its shape and strides describe every stored element without
30//! hidden gaps.
31//!
32//! The allocation start is aligned to 64 bytes by [`Vec64`], allowing efficient
33//! SIMD operations over the flattened buffer.
34//!
35//! Per-column padding for CPU-oriented BLAS and LAPACK access is provided by
36//! [`Matrix`] instead. Converting an `NdArray` with `to_matrix` re-lays the data
37//! into that padded representation when required.
38//!
39//! NumPy and PyTorch can consume the column-major strides directly. Consumers
40//! that require row-major C-contiguous storage must create a contiguous copy,
41//! such as with `np.ascontiguousarray(...)` or `.contiguous()`.
42//!
43//! ## Null values
44//!
45//! `NdArray` does not store a null mask. When converting nullable floating-point
46//! data, nulls are represented as `NaN`. The resulting array cannot distinguish
47//! a source null from an ordinary `NaN`; interpretation is left to the consuming
48//! kernel or framework.
49//!
50//! ## Interoperability
51//!
52//! - [`Matrix`] to a two-dimensional `NdArray`: zero-copy for `f64`; the padded
53//!   column stride is preserved.
54//! - Two-dimensional `NdArray` to [`Matrix`]: re-lays data into the padded column
55//!   format for `f64`, or transfers without copying when the existing stride is
56//!   already compatible.
57//! - Two-dimensional `NdArray` to [`Table`]: copies each `f64` column into a
58//!   separate 64-byte-aligned `FloatArray`.
59//! - One-dimensional `NdArray` to [`Array`]: moves the `f64` buffer into a
60//!   `FloatArray<f64>`.
61//! - [`Table`] to `NdArray` through [`TryFrom`]: copies `f64` values and converts
62//!   nulls to `NaN`.
63//! - DLPack import and export: shares data with compatible `f32` and `f64`
64//!   consumers. Whether a transfer is zero-copy depends on ownership, alignment,
65//!   layout, and protocol constraints.
66
67use std::fmt;
68use std::ops::{Index, IndexMut, Range, RangeFrom, RangeFull, RangeTo};
69use std::sync::Arc;
70
71use crate::enums::error::MinarrowError;
72use crate::enums::shape_dim::ShapeDim;
73use crate::structs::buffer::Buffer;
74#[cfg(all(feature = "views", feature = "select"))]
75use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
76use crate::traits::type_unions::Float;
77use crate::traits::{concatenate::Concatenate, shape::Shape};
78use crate::{Array, ArrowType, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
79
80#[cfg(feature = "matrix")]
81use crate::structs::matrix::{Matrix, aligned_stride};
82#[cfg(feature = "views")]
83use crate::structs::views::ndarray_view::NdArrayV;
84#[cfg(feature = "dlpack")]
85use crate::ffi::dlpack::{
86    export_to_dlpack, export_to_dlpack_versioned, DLPackTensor, DLPackTensorVersioned,
87};
88
89// ****************************************************************
90// NdArray
91// ****************************************************************
92
93/// N-dimensional dense array of float values.
94///
95/// Backed by [`Buffer<T>`] for zero-copy interop with external memory.
96/// Compact column-major layout with a 64-byte aligned allocation start.
97///
98/// See the [module-level documentation](self) for design rationale.
99#[derive(Clone)]
100pub struct NdArray<T> {
101    pub(crate) data: Arc<Buffer<T>>,
102    pub(crate) dims: NdDims,
103    pub name: Option<String>,
104}
105
106/// Logical equality over shape and values in logical order. Strides and
107/// the name do not affect equality, matching the view and SuperNdArray types.
108impl<T: Float> PartialEq for NdArray<T> {
109    fn eq(&self, other: &Self) -> bool {
110        if self.dims.shape() != other.dims.shape() {
111            return false;
112        }
113        self.into_iter()
114            .zip(other.into_iter())
115            .all(|(a, b)| a == b)
116    }
117}
118
119// *** Construction ************************************************
120
121impl<T: Float> NdArray<T> {
122    /// Create a zeroed NdArray with the given shape, column-major strides.
123    /// `shape == &[]` creates a rank-zero array containing one scalar value;
124    /// `shape == &[0]` creates an empty rank-one array.
125    pub fn new(shape: &[usize]) -> Self {
126        let dims = NdDims::from_shape(shape);
127        let total = buffer_len(dims.shape(), dims.strides());
128        let mut v = Vec64::with_capacity(total);
129        v.0.resize(total, T::default());
130        NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
131    }
132
133    /// Create a zeroed NdArray with a name.
134    pub fn new_named(shape: &[usize], name: impl Into<String>) -> Self {
135        let mut arr = Self::new(shape);
136        arr.name = Some(name.into());
137        arr
138    }
139
140    /// Create from a flat column-major slice and shape.
141    ///
142    /// The slice holds `product(shape)` logical elements in column-major
143    /// order, matching the compact layout, so the data copies straight in.
144    /// A shape with zero axes (`&[]`) has product one and therefore requires
145    /// one value. A shape containing a zero-length axis, such as `&[0]`, has
146    /// product zero and requires no values.
147    pub fn from_slice(data: &[T], shape: &[usize]) -> Self {
148        let logical_len: usize = shape.iter().product();
149        assert_eq!(
150            data.len(), logical_len,
151            "NdArray::from_slice: data length {} does not match shape product {}",
152            data.len(), logical_len
153        );
154        let dims = NdDims::from_shape(shape);
155        NdArray {
156            data: Arc::new(Buffer::from_slice(data)),
157            dims,
158            name: None,
159        }
160    }
161
162    /// Create from an owned 64-byte aligned vector, moving it in without
163    /// a copy. The vector holds `product(shape)` logical elements in
164    /// column-major order, matching the compact layout. A shape with zero
165    /// axes (`&[]`) requires one value, while `&[0]` requires none.
166    pub fn from_vec64(data: Vec64<T>, shape: &[usize]) -> Self {
167        let logical_len: usize = shape.iter().product();
168        assert_eq!(
169            data.len(), logical_len,
170            "NdArray::from_vec64: data length {} does not match shape product {}",
171            data.len(), logical_len
172        );
173        let dims = NdDims::from_shape(shape);
174        NdArray {
175            data: Arc::new(Buffer::from_vec64(data)),
176            dims,
177            name: None,
178        }
179    }
180
181    /// Create from a pre-owned `Buffer<f64>` with explicit shape and strides.
182    ///
183    /// The buffer must already contain `buffer_len(shape, strides)` elements
184    /// in the correct strided layout.
185    pub fn from_buffer(data: Buffer<T>, shape: &[usize], strides: &[usize]) -> Self {
186        let required = buffer_len(shape, strides);
187        assert!(
188            data.len() >= required,
189            "NdArray::from_buffer: buffer has {} elements but shape requires {}",
190            data.len(), required
191        );
192        let dims = NdDims::from_shape_and_strides(shape, strides);
193        NdArray { data: Arc::new(data), dims, name: None }
194    }
195
196    /// Create an NdArray filled with a constant value.
197    pub fn fill(shape: &[usize], value: T) -> Self {
198        let dims = NdDims::from_shape(shape);
199        let total = buffer_len(dims.shape(), dims.strides());
200        let mut v = Vec64::with_capacity(total);
201        v.0.resize(total, value);
202        NdArray { data: Arc::new(Buffer::from_vec64(v)), dims, name: None }
203    }
204
205    /// Create an NdArray of ones.
206    pub fn ones(shape: &[usize]) -> Self {
207        Self::fill(shape, T::one())
208    }
209
210    /// Create a 2D identity matrix.
211    pub fn eye(n: usize) -> Self {
212        let mut arr = Self::new(&[n, n]);
213        let stride = arr.dims.strides()[1];
214        let buf = Arc::make_mut(&mut arr.data).as_mut_slice();
215        for i in 0..n {
216            buf[i * stride + i] = T::one();
217        }
218        arr
219    }
220
221    /// Create a 1D NdArray with evenly spaced values in `[start, end]`.
222    pub fn linspace(start: T, end: T, n: usize) -> Self {
223        assert!(n >= 2, "linspace requires at least 2 points");
224        let step = (end - start) / T::from(n - 1).unwrap();
225        let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
226        NdArray {
227            data: Arc::new(Buffer::from_vec64(v)),
228            dims: NdDims::from_shape(&[n]),
229            name: None,
230        }
231    }
232
233    /// Create a 1D NdArray with values `start, start+step, start+2*step, ...`
234    /// for `n` elements.
235    pub fn arange(start: T, step: T, n: usize) -> Self {
236        let v: Vec64<T> = (0..n).map(|i| start + step * T::from(i).unwrap()).collect();
237        NdArray {
238            data: Arc::new(Buffer::from_vec64(v)),
239            dims: NdDims::from_shape(&[n]),
240            name: None,
241        }
242    }
243
244    // *** Shape and introspection *********************************
245
246    /// Number of dimensions.
247    #[inline]
248    pub fn ndim(&self) -> usize { self.dims.ndim() }
249
250    /// Shape as a slice of dimension sizes.
251    #[inline]
252    pub fn shape(&self) -> &[usize] { self.dims.shape() }
253
254    /// Strides as a slice of element offsets per dimension.
255    #[inline]
256    pub fn strides(&self) -> &[usize] { self.dims.strides() }
257
258    /// Total number of logical elements i.e. the product of shape.
259    #[inline]
260    pub fn len(&self) -> usize { self.dims.len() }
261
262    /// Leading-axis (axis 0) observation count i.e. `shape()[0]`,
263    /// matching `SuperNdArray::n_obs` and NumPy's `len(arr)`.
264    /// Panics for a rank-zero scalar, which has no axis 0.
265    #[inline]
266    pub fn n_obs(&self) -> usize {
267        assert!(self.ndim() > 0, "n_obs() requires an axis 0");
268        self.dims.shape()[0]
269    }
270
271    /// True if any dimension is zero.
272    #[inline]
273    pub fn is_empty(&self) -> bool { self.len() == 0 }
274
275    /// True if the buffer layout matches compact column-major strides,
276    /// with no transposition or non-standard stride pattern.
277    pub fn is_contiguous(&self) -> bool {
278        let shape = self.dims.shape();
279        let strides = self.dims.strides();
280        let mut expected = 1;
281        for d in 0..shape.len() {
282            if strides[d] != expected {
283                return false;
284            }
285            expected *= shape[d];
286        }
287        true
288    }
289
290    /// True if any element is NaN.
291    pub fn has_nan(&self) -> bool {
292        self.into_iter().any(|v| v.is_nan())
293    }
294
295    /// Count of NaN elements.
296    pub fn nan_count(&self) -> usize {
297        self.into_iter().filter(|v| v.is_nan()).count()
298    }
299
300    // *** Element access ******************************************
301
302    /// Get element by N-dimensional index. Panics if out of bounds.
303    #[inline]
304    pub fn get(&self, indices: &[usize]) -> T {
305        self.data.as_slice()[self.offset_of(indices)]
306    }
307
308    /// Like `get`, but skips bounds checks.
309    ///
310    /// # Safety
311    /// The caller guarantees each index is within its dimension.
312    #[inline(always)]
313    pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
314        let strides = self.dims.strides();
315        let mut off = 0;
316        for d in 0..indices.len() {
317            off += indices[d] * strides[d];
318        }
319        // SAFETY: in-bounds indices produce an in-bounds flat offset.
320        unsafe { *self.data.as_slice().get_unchecked(off) }
321    }
322
323    /// Set element by N-dimensional index. Panics if out of bounds.
324    /// Triggers copy-on-write when views share the buffer.
325    ///
326    /// For repeated writes, detach once and write through the slice:
327    /// ```ignore
328    /// let s1 = a.strides()[1];
329    /// let s = a.as_mut_slice();      // one copy-on-write detach
330    /// for (i, j, v) in writes {
331    ///     // SAFETY: i and j are within shape
332    ///     unsafe { *s.get_unchecked_mut(i + j * s1) = v; }
333    /// }
334    /// ```
335    #[inline]
336    pub fn set(&mut self, indices: &[usize], value: T) {
337        let off = self.offset_of(indices);
338        Arc::make_mut(&mut self.data).as_mut_slice()[off] = value;
339    }
340
341    /// Compute flat buffer offset for an N-dimensional index.
342    #[inline]
343    pub(crate) fn offset_of(&self, indices: &[usize]) -> usize {
344        offset_of_impl(indices, self.dims.shape(), self.dims.strides())
345    }
346
347    // *** Metadata ************************************************
348
349    /// Set the array name.
350    #[inline]
351    pub fn set_name(&mut self, name: impl Into<String>) {
352        self.name = Some(name.into());
353    }
354
355    /// Immutable reference to the full flat buffer.
356    #[inline]
357    pub fn as_slice(&self) -> &[T] {
358        self.data.as_slice()
359    }
360
361    /// Mutable reference to the full flat buffer. Triggers copy-on-write
362    /// when views share the buffer.
363    #[inline]
364    pub fn as_mut_slice(&mut self) -> &mut [T] {
365        Arc::make_mut(&mut self.data).as_mut_slice()
366    }
367
368    /// Fill every logical element with a value.
369    pub fn fill_with(&mut self, value: T) {
370        // For contiguous arrays, fill the whole buffer
371        if self.is_contiguous() {
372            Arc::make_mut(&mut self.data).as_mut_slice().fill(value);
373            return;
374        }
375        // Walk logical positions for non-contiguous layouts
376        let offsets: Vec<usize> = {
377            let shape = self.dims.shape();
378            let strides = self.dims.strides();
379            let mut result = Vec::with_capacity(self.len());
380            let mut indices = vec![0usize; shape.len()];
381            for _ in 0..self.len() {
382                result.push(indices.iter().zip(strides).map(|(&i, &s)| i * s).sum());
383                let mut carry = true;
384                for d in 0..shape.len() {
385                    if carry {
386                        indices[d] += 1;
387                        if indices[d] < shape[d] { carry = false; } else { indices[d] = 0; }
388                    }
389                }
390            }
391            result
392        };
393        let buf = Arc::make_mut(&mut self.data).as_mut_slice();
394        for off in offsets { buf[off] = value; }
395    }
396
397    // *** 2D axis access (column-major) ***************************
398
399    /// Immutable column slice for a 2D array.
400    /// Panics if ndim != 2, the column index is out of bounds, or axis 0
401    /// is not unit-stride i.e. column elements are not contiguous.
402    #[inline]
403    pub fn col(&self, col: usize) -> &[T] {
404        let shape = self.dims.shape();
405        assert_eq!(shape.len(), 2, "col() requires a 2D array");
406        assert!(col < shape[1], "Column index out of bounds");
407        assert_eq!(
408            self.dims.strides()[0], 1,
409            "col() requires unit stride on axis 0; call to_contiguous() first"
410        );
411        let stride = self.dims.strides()[1];
412        let start = col * stride;
413        &self.data.as_slice()[start..start + shape[0]]
414    }
415
416    /// Mutable column slice for a 2D array.
417    /// Panics if ndim != 2, the column index is out of bounds, or axis 0
418    /// is not unit-stride i.e. column elements are not contiguous.
419    /// Triggers copy-on-write if the buffer is shared.
420    #[inline]
421    pub fn col_mut(&mut self, col: usize) -> &mut [T] {
422        let shape = self.dims.shape();
423        assert_eq!(shape.len(), 2, "col_mut() requires a 2D array");
424        assert!(col < shape[1], "Column index out of bounds");
425        assert_eq!(
426            self.dims.strides()[0], 1,
427            "col_mut() requires unit stride on axis 0; call to_contiguous() first"
428        );
429        let stride = self.dims.strides()[1];
430        let n_rows = shape[0];
431        let start = col * stride;
432        &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
433    }
434
435    /// All columns as immutable slices. 2D only.
436    /// Panics if ndim != 2 or axis 0 is not unit-stride.
437    pub fn columns(&self) -> Vec<&[T]> {
438        let shape = self.dims.shape();
439        assert_eq!(shape.len(), 2, "columns() requires a 2D array");
440        assert_eq!(
441            self.dims.strides()[0], 1,
442            "columns() requires unit stride on axis 0; call to_contiguous() first"
443        );
444        let stride = self.dims.strides()[1];
445        let n_rows = shape[0];
446        let buf = self.data.as_slice();
447        (0..shape[1])
448            .map(|c| &buf[c * stride..c * stride + n_rows])
449            .collect()
450    }
451
452    /// All columns as mutable slices. 2D only.
453    /// Panics if ndim != 2 or axis 0 is not unit-stride.
454    pub fn columns_mut(&mut self) -> Vec<&mut [T]> {
455        let shape = self.dims.shape();
456        assert_eq!(shape.len(), 2, "columns_mut() requires a 2D array");
457        assert_eq!(
458            self.dims.strides()[0], 1,
459            "columns_mut() requires unit stride on axis 0; call to_contiguous() first"
460        );
461        let n_rows = shape[0];
462        let n_cols = shape[1];
463        let stride = self.dims.strides()[1];
464        let ptr = Arc::make_mut(&mut self.data).as_mut_slice().as_mut_ptr();
465        let mut result = Vec::with_capacity(n_cols);
466        for c in 0..n_cols {
467            let start = c * stride;
468            // SAFETY: each slice is within bounds and non-overlapping,
469            // we have exclusive &mut access.
470            unsafe {
471                let col_ptr = ptr.add(start);
472                result.push(std::slice::from_raw_parts_mut(col_ptr, n_rows));
473            }
474        }
475        result
476    }
477
478    /// Wrap as a full `NdArrayV` view for repeated observation access.
479    ///
480    /// Call `.obs(i)` on the returned view to get individual observations.
481    /// The view holds the parent through the array's shared internal
482    /// buffer, so this is a refcount bump.
483    #[cfg(feature = "views")]
484    pub fn as_view(&self) -> NdArrayV<T> {
485        NdArrayV::from_ndarray(self.clone())
486    }
487
488    /// Zero-copy view of a single observation (axis-0 element).
489    ///
490    /// Returns an (N-1)-dimensional `NdArrayV` view. For a 2D array
491    /// with shape `[n, m]`, returns a 1D view of shape `[m]`. For 3D
492    /// `[n, m, k]`, returns 2D `[m, k]`. Requires rank 2 or higher - for
493    /// scalar access on a 1D array use `get(&[i])`.
494    ///
495    /// For repeated access in a loop, prefer `nd.as_view()` then call
496    /// `.obs()` on the view to avoid re-wrapping each time.
497    #[cfg(feature = "views")]
498    pub fn obs(&self, idx: usize) -> NdArrayV<T> {
499        self.as_view().obs(idx)
500    }
501
502    // *** BLAS/LAPACK compatibility (2D) **************************
503
504    /// Number of rows as i32. Panics if ndim != 2.
505    #[inline]
506    pub fn m(&self) -> i32 {
507        assert_eq!(self.ndim(), 2, "m() requires a 2D array");
508        self.dims.shape()[0] as i32
509    }
510
511    /// Number of columns as i32. Panics if ndim != 2.
512    #[inline]
513    pub fn n(&self) -> i32 {
514        assert_eq!(self.ndim(), 2, "n() requires a 2D array");
515        self.dims.shape()[1] as i32
516    }
517
518    /// Leading dimension for BLAS. Panics if ndim != 2.
519    #[inline]
520    pub fn lda(&self) -> i32 {
521        assert_eq!(self.ndim(), 2, "lda() requires a 2D array");
522        self.dims.strides()[1] as i32
523    }
524
525    // *** Reshape *************************************************
526
527    /// Reshape to a new shape. Returns a new NdArray with re-laid out data.
528    ///
529    /// The total number of logical elements must match. Data is copied
530    /// from logical element order into the new strided layout.
531    pub fn reshape(&self, new_shape: &[usize]) -> Result<NdArray<T>, MinarrowError> {
532        let new_len: usize = new_shape.iter().product();
533        if new_len != self.len() {
534            return Err(MinarrowError::ShapeError {
535                message: format!(
536                    "reshape: cannot reshape array of size {} into shape {:?}",
537                    self.len(), new_shape
538                ),
539            });
540        }
541        let flat: Vec64<T> = self.into_iter().collect();
542        let mut result = NdArray::from_slice(&flat, new_shape);
543        result.name = self.name.clone();
544        Ok(result)
545    }
546
547    /// Transpose by reversing the axis order. Returns a new NdArray with
548    /// re-laid out data. A 1D array copies through unchanged.
549    ///
550    /// For a zero-copy transposed view, call `as_view()` and use the
551    /// view's `transpose()`.
552    pub fn transpose(&self) -> NdArray<T> {
553        let shape = self.dims.shape();
554        let ndim = shape.len();
555        if ndim <= 1 {
556            let mut result = self.to_contiguous();
557            result.name = self.name.clone();
558            return result;
559        }
560        if ndim == 2 {
561            let (n_rows, n_cols) = (shape[0], shape[1]);
562            let new_shape = [n_cols, n_rows];
563            let mut result = NdArray::new(&new_shape);
564            let src_stride = self.dims.strides()[1];
565            let dst_stride = result.dims.strides()[1];
566            let src = self.data.as_slice();
567            let dst = Arc::make_mut(&mut result.data).as_mut_slice();
568            for r in 0..n_rows {
569                for c in 0..n_cols {
570                    dst[r * dst_stride + c] = src[c * src_stride + r];
571                }
572            }
573            result.name = self.name.clone();
574            return result;
575        }
576
577        // General N-D. Walking the result's logical positions in column-major
578        // order reads the source at the reversed index, which is the source
579        // walked with reversed strides.
580        let new_shape: Vec<usize> = shape.iter().rev().copied().collect();
581        let rev_strides: Vec<usize> = self.dims.strides().iter().rev().copied().collect();
582        let src = self.data.as_slice();
583        let total = self.len();
584        let mut buf = Vec64::with_capacity(total);
585        let mut indices = vec![0usize; ndim];
586        for _ in 0..total {
587            let offset: usize = indices.iter()
588                .zip(rev_strides.iter())
589                .map(|(&i, &s)| i * s)
590                .sum();
591            buf.push(src[offset]);
592            let mut carry = true;
593            for d in 0..ndim {
594                if carry {
595                    indices[d] += 1;
596                    if indices[d] < new_shape[d] {
597                        carry = false;
598                    } else {
599                        indices[d] = 0;
600                    }
601                }
602            }
603        }
604        NdArray {
605            data: Arc::new(Buffer::from_vec64(buf)),
606            dims: NdDims::from_shape(&new_shape),
607            name: self.name.clone(),
608        }
609    }
610
611    /// Flatten to a contiguous 1D array.
612    pub fn flatten(&self) -> NdArray<T> {
613        let flat: Vec64<T> = self.into_iter().collect();
614        let n = flat.len();
615        NdArray {
616            data: Arc::new(Buffer::from_vec64(flat)),
617            dims: NdDims::from_shape(&[n]),
618            name: self.name.clone(),
619        }
620    }
621
622    /// If the array has non-standard strides, re-lay out into default
623    /// column-major contiguous form.
624    pub fn to_contiguous(&self) -> NdArray<T> {
625        if self.is_contiguous() {
626            return self.clone();
627        }
628        let mut result = NdArray::from_slice(
629            &self.into_iter().collect::<Vec64<T>>(),
630            self.shape(),
631        );
632        result.name = self.name.clone();
633        result
634    }
635
636    // *** Slicing: arr.slice(nd![1..4, 2..5]) *************************
637
638    /// Slice this array along any combination of axes.
639    ///
640    /// Each axis takes any [`DataSelector`] - a single index collapses
641    /// that dimension, and a contiguous range keeps it. Returns a
642    /// zero-copy `NdArrayV` view that holds the parent through the
643    /// array's shared internal buffer, so this is a refcount bump,
644    /// matching `as_view` and `obs`.
645    ///
646    /// # Examples
647    /// ```ignore
648    /// arr.slice(&[&2])              // single index on 1D
649    /// arr.slice(&[&(1..4)])         // range on axis 0
650    /// arr.slice(nd![1..4, 2..5])    // range on both axes (2D)
651    /// arr.slice(nd![0..3, 2])       // range on axis 0, single on axis 1
652    /// arr.slice(nd![1, 0..4, 3])    // mixed for 3D
653    /// ```
654    #[cfg(all(feature = "views", feature = "select"))]
655    pub fn slice(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
656        assert_eq!(
657            selection.len(), self.ndim(),
658            "slice(): expected {} axes, got {}", self.ndim(), selection.len()
659        );
660
661        let shape = self.dims.shape();
662        let strides = self.dims.strides();
663
664        // Compute the new offset, shape, and strides
665        let mut new_offset: usize = 0;
666        let mut new_shape = Vec::with_capacity(self.ndim());
667        let mut new_strides = Vec::with_capacity(self.ndim());
668
669        for (d, sel) in selection.iter().enumerate() {
670            let (start, end, collapse) = sel.resolve_axis(shape[d]);
671            assert!(
672                end <= shape[d],
673                "slice(): end {} out of bounds for axis {} (size {})", end, d, shape[d]
674            );
675            new_offset += start * strides[d];
676            if !collapse {
677                new_shape.push(end - start);
678                new_strides.push(strides[d]);
679            }
680        }
681
682        NdArrayV::new(self.clone(), new_offset, &new_shape, &new_strides)
683    }
684
685    // *** Apply ***************************************************
686
687    /// Apply a function to every logical element, returning a new compact
688    /// array with this array's shape and name. The caller supplies the
689    /// operation; NdArray supplies logical-order traversal and materialisation.
690    pub fn apply(&self, f: impl Fn(T) -> T) -> NdArray<T> {
691        let flat: Vec64<T> = self.into_iter().map(f).collect();
692        let mut result = NdArray::from_slice(&flat, self.shape());
693        result.name = self.name.clone();
694        result
695    }
696
697    /// Apply a function to every logical element in place, with no
698    /// allocation. Copy-on-write triggers first when views share the
699    /// buffer.
700    pub fn apply_mut(&mut self, f: impl Fn(T) -> T) {
701        if self.is_contiguous() {
702            for v in Arc::make_mut(&mut self.data).as_mut_slice() {
703                *v = f(*v);
704            }
705            return;
706        }
707        // Walk logical positions for non-contiguous layouts so stride
708        // padding stays untouched.
709        let shape = self.dims.shape().to_vec();
710        let strides = self.dims.strides().to_vec();
711        let total = self.len();
712        let buf = Arc::make_mut(&mut self.data).as_mut_slice();
713        let mut indices = vec![0usize; shape.len()];
714        for _ in 0..total {
715            let offset: usize = indices.iter().zip(strides.iter()).map(|(&i, &s)| i * s).sum();
716            buf[offset] = f(buf[offset]);
717            let mut carry = true;
718            for d in 0..shape.len() {
719                if carry {
720                    indices[d] += 1;
721                    if indices[d] < shape[d] {
722                        carry = false;
723                    } else {
724                        indices[d] = 0;
725                    }
726                }
727            }
728        }
729    }
730
731    /// Apply a function to every 1D lane along the given axis, collapsing
732    /// that axis. Each lane arrives as a zero-copy [`NdArrayV`] and the
733    /// closure returns one value for it, so the output shape drops `axis`.
734    /// Requires rank 2 or higher - a 1D array is itself a single lane.
735    #[cfg(feature = "views")]
736    pub fn apply_axis(&self, axis: usize, mut f: impl FnMut(NdArrayV<T>) -> T) -> NdArray<T> {
737        let shape = self.dims.shape();
738        let strides = self.dims.strides();
739        let ndim = shape.len();
740        assert!(ndim >= 2, "apply_axis requires a 2D or higher array");
741        assert!(axis < ndim, "apply_axis: axis {} out of bounds for {}D array", axis, ndim);
742
743        let out_shape: Vec<usize> = shape
744            .iter()
745            .enumerate()
746            .filter(|(d, _)| *d != axis)
747            .map(|(_, &s)| s)
748            .collect();
749        let out_dims: Vec<usize> = (0..ndim).filter(|&d| d != axis).collect();
750        let total: usize = out_shape.iter().product();
751
752        let lane_shape = [shape[axis]];
753        let lane_strides = [strides[axis]];
754
755        // Walk the output positions in column-major order. Each position
756        // anchors one lane's base offset in the source.
757        let mut flat = Vec64::with_capacity(total);
758        let mut indices = vec![0usize; out_shape.len()];
759        for _ in 0..total {
760            let offset: usize = indices
761                .iter()
762                .zip(out_dims.iter())
763                .map(|(&i, &d)| i * strides[d])
764                .sum();
765            let lane = NdArrayV::new(self.clone(), offset, &lane_shape, &lane_strides);
766            flat.push(f(lane));
767            let mut carry = true;
768            for d in 0..out_shape.len() {
769                if carry {
770                    indices[d] += 1;
771                    if indices[d] < out_shape[d] {
772                        carry = false;
773                    } else {
774                        indices[d] = 0;
775                    }
776                }
777            }
778        }
779        let mut result = NdArray::from_slice(&flat, &out_shape);
780        result.name = self.name.clone();
781        result
782    }
783
784    // *** Conversions *********************************************
785
786    /// Export as a legacy DLPack tensor for PyTorch, NumPy, JAX, and other
787    /// compatible consumers. Shared or aliased storage copies because the
788    /// legacy protocol cannot mark the exported tensor read-only.
789    ///
790    /// Returns a `DLPackTensor` that manages the lifecycle. Drop it to
791    /// release, or call `.into_raw()` to transfer ownership to an FFI
792    /// consumer such as a PyCapsule.
793    #[cfg(feature = "dlpack")]
794    pub fn to_dlpack(self) -> DLPackTensor {
795        export_to_dlpack(self)
796    }
797
798    /// Export as a DLPack 1.x versioned tensor, carrying the spec version
799    /// and flags fields that PyTorch and JAX negotiate. The read-only
800    /// flag is set when the backing buffer is shared.
801    #[cfg(feature = "dlpack")]
802    pub fn to_dlpack_versioned(self) -> DLPackTensorVersioned {
803        export_to_dlpack_versioned(self)
804    }
805
806    // *** Parallel iteration (rayon) ******************************
807
808    /// Parallel iterator over the underlying buffer. Rayon splits
809    /// the contiguous data into chunks across threads automatically.
810    #[cfg(feature = "parallel_proc")]
811    pub fn par_iter(&self) -> rayon::slice::Iter<'_, T>
812    where
813        T: Send + Sync,
814    {
815        use rayon::prelude::*;
816        self.data.as_slice().par_iter()
817    }
818
819    /// Parallel chunks of the underlying buffer. Each chunk is a
820    /// contiguous `&[T]` slice that rayon distributes across threads.
821    #[cfg(feature = "parallel_proc")]
822    pub fn par_chunks(&self, chunk_size: usize) -> rayon::slice::Chunks<'_, T>
823    where
824        T: Send + Sync,
825    {
826        use rayon::prelude::*;
827        self.data.as_slice().par_chunks(chunk_size)
828    }
829
830    /// Iterator over axis-0 observations. Each item is the observation
831    /// index and a zero-copy `NdArrayV` view.
832    #[cfg(feature = "views")]
833    pub fn iter_obs(&self) -> impl Iterator<Item = (usize, NdArrayV<T>)> + '_ {
834        assert!(self.ndim() >= 2, "iter_obs() requires a 2D or higher array");
835        let n_obs = self.shape()[0];
836        (0..n_obs).map(move |i| (i, self.obs(i)))
837    }
838
839    /// Parallel iterator over axis-0 observations. Each item is the
840    /// observation index and a zero-copy `NdArrayV` view.
841    #[cfg(all(feature = "parallel_proc", feature = "views"))]
842    pub fn par_iter_obs(&self) -> impl rayon::iter::ParallelIterator<Item = (usize, NdArrayV<T>)> + '_
843    where
844        T: Send + Sync,
845    {
846        use rayon::prelude::*;
847        assert!(self.ndim() >= 2, "par_iter_obs() requires a 2D or higher array");
848        let n_obs = self.dims.shape()[0];
849        (0..n_obs).into_par_iter().map(move |i| (i, self.obs(i)))
850    }
851
852    /// Iterate one logical axis-0 run identified by its flattened outer
853    /// index. This composes the logical iterator for SuperNdArray without
854    /// materialising its batches.
855    pub(crate) fn iter_axis0_run(&self, run_idx: usize) -> impl ExactSizeIterator<Item = T> + '_ {
856        assert!(self.ndim() > 0, "axis-0 iteration requires an axis 0");
857        let n_runs: usize = self.shape()[1..].iter().product();
858        assert!(run_idx < n_runs, "axis-0 run {} out of bounds ({})", run_idx, n_runs);
859
860        let mut rem = run_idx;
861        let mut offset = 0;
862        for d in 1..self.ndim() {
863            offset += (rem % self.shape()[d]) * self.strides()[d];
864            rem /= self.shape()[d];
865        }
866        let stride = self.strides()[0];
867        (0..self.shape()[0]).map(move |i| self.data.as_slice()[offset + i * stride])
868    }
869}
870
871// *** f64 conversions to the Array/Table/Matrix enum boundary *****
872
873impl NdArray<f64> {
874    /// Convert a 2D NdArray to a Table.
875    ///
876    /// Each column is copied into its own 64-byte aligned `FloatArray<f64>`,
877    /// since the compact tensor layout does not place column starts on
878    /// alignment boundaries. `fields` must have exactly `n_cols` entries
879    /// when supplied. Each supplied field must describe a non-nullable
880    /// [`ArrowType::Float64`] column. `None` generates `col_0`, `col_1`,
881    /// and so on.
882    pub fn to_table(self, fields: Option<Vec<Field>>) -> Result<Table, MinarrowError> {
883        let shape = self.dims.shape();
884        if shape.len() != 2 {
885            return Err(MinarrowError::ShapeError {
886                message: format!("to_table requires a 2D array, got {}D", shape.len()),
887            });
888        }
889        let n_cols = shape[1];
890        let fields = match fields {
891            Some(fields) => {
892                if fields.len() != n_cols {
893                    return Err(MinarrowError::ShapeError {
894                        message: format!(
895                            "to_table: expected {} fields for {} columns, got {}",
896                            n_cols, n_cols, fields.len()
897                        ),
898                    });
899                }
900                for (index, field) in fields.iter().enumerate() {
901                    if field.dtype != ArrowType::Float64 {
902                        return Err(MinarrowError::TypeError {
903                            from: "Field",
904                            to: "Float64 NdArray column",
905                            message: Some(format!(
906                                "to_table: field {} ('{}') has dtype {:?}, expected Float64",
907                                index, field.name, field.dtype
908                            )),
909                        });
910                    }
911                    if field.nullable {
912                        return Err(MinarrowError::NullError {
913                            message: Some(format!(
914                                "to_table: field {} ('{}') is nullable, but NdArray columns are non-nullable",
915                                index, field.name
916                            )),
917                        });
918                    }
919                }
920                fields
921            }
922            None => (0..n_cols)
923                .map(|i| Field::new(format!("col_{}", i), ArrowType::Float64, false, None))
924                .collect(),
925        };
926        // The column copies below read unit-stride axis-0 runs, so a
927        // transposed or otherwise strided layout re-lays first.
928        let this = if self.dims.strides()[0] == 1 { self } else { self.to_contiguous() };
929        let n_rows = this.dims.shape()[0];
930        let stride = this.dims.strides()[1];
931        let name = this.name;
932        let buf = this.data.as_slice();
933
934        let mut cols = Vec::with_capacity(n_cols);
935        for (i, field) in fields.into_iter().enumerate() {
936            let col_start = i * stride;
937            let col: Buffer<f64> = Buffer::from_slice(&buf[col_start..col_start + n_rows]);
938            let float_arr = FloatArray::new(col, None);
939            let array = Array::NumericArray(NumericArray::Float64(Arc::new(float_arr)));
940            cols.push(FieldArray::new(field, array));
941        }
942
943        Ok(Table::new(name.unwrap_or_default(), Some(cols)))
944    }
945
946    /// Convert a 1D NdArray to an Array (FloatArray<f64>).
947    /// Compact, exact-length storage moves across directly. A strided array
948    /// or one with unused backing elements is materialised in logical order.
949    pub fn to_array(self) -> Result<Array, MinarrowError> {
950        if self.ndim() != 1 {
951            return Err(MinarrowError::ShapeError {
952                message: format!("to_array requires a 1D array, got {}D", self.ndim()),
953            });
954        }
955        let buffer = if self.dims.strides()[0] == 1 && self.data.len() == self.len() {
956            Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone())
957        } else {
958            Buffer::from_vec64((&self).into_iter().collect())
959        };
960        let float_arr = FloatArray::new(buffer, None);
961        Ok(Array::NumericArray(NumericArray::Float64(Arc::new(float_arr))))
962    }
963
964    /// Convert a 2D NdArray to a Matrix.
965    ///
966    /// Matrix pads each column to a 64-byte boundary for BLAS/LAPACK and
967    /// SIMD access, so compact tensor data is re-laid out into the padded
968    /// form. An array whose stride already matches the padded layout, such
969    /// as one built from a Matrix, moves across zero-copy.
970    #[cfg(feature = "matrix")]
971    pub fn to_matrix(self) -> Result<Matrix, MinarrowError> {
972        let shape = self.dims.shape();
973        if shape.len() != 2 {
974            return Err(MinarrowError::ShapeError {
975                message: format!("to_matrix requires a 2D array, got {}D", shape.len()),
976            });
977        }
978        let n_rows = shape[0];
979        let n_cols = shape[1];
980        let strides = self.dims.strides();
981
982        if strides[0] == 1 && strides[1] == aligned_stride(n_rows) {
983            return Ok(Matrix {
984                n_rows,
985                n_cols,
986                stride: strides[1],
987                data: Arc::try_unwrap(self.data).unwrap_or_else(|arc| (*arc).clone()),
988                name: self.name,
989            });
990        }
991
992        let name = self.name.clone();
993        let compact: Vec64<f64> = self.into_iter().collect();
994        Ok(Matrix::from_f64_unaligned(&compact, n_rows, n_cols, name))
995    }
996}
997
998// ****************************************************************
999// NdDims - internal dimension storage
1000// ****************************************************************
1001
1002/// Internal storage for shape and strides.
1003///
1004/// Inline arrays for 1D-5D to avoid heap allocation; the `Dn` variant
1005/// handles 6+ dimensions via boxed slices.
1006#[derive(Clone, PartialEq)]
1007pub(crate) enum NdDims {
1008    D1 { shape: [usize; 1], strides: [usize; 1] },
1009    D2 { shape: [usize; 2], strides: [usize; 2] },
1010    D3 { shape: [usize; 3], strides: [usize; 3] },
1011    D4 { shape: [usize; 4], strides: [usize; 4] },
1012    D5 { shape: [usize; 5], strides: [usize; 5] },
1013    Dn { shape: Box<[usize]>, strides: Box<[usize]> },
1014}
1015
1016impl NdDims {
1017    /// Build dims from a shape slice, computing compact column-major
1018    /// strides.
1019    pub(crate) fn from_shape(shape: &[usize]) -> Self {
1020        let strides = col_major_strides(shape);
1021        Self::from_shape_and_strides(shape, &strides)
1022    }
1023
1024    /// Build dims from explicit shape and strides.
1025    pub(crate) fn from_shape_and_strides(shape: &[usize], strides: &[usize]) -> Self {
1026        assert_eq!(
1027            shape.len(),
1028            strides.len(),
1029            "NdArray: shape rank {} does not match strides rank {}",
1030            shape.len(),
1031            strides.len()
1032        );
1033        match shape.len() {
1034            1 => NdDims::D1 {
1035                shape: [shape[0]],
1036                strides: [strides[0]],
1037            },
1038            2 => NdDims::D2 {
1039                shape: [shape[0], shape[1]],
1040                strides: [strides[0], strides[1]],
1041            },
1042            3 => NdDims::D3 {
1043                shape: [shape[0], shape[1], shape[2]],
1044                strides: [strides[0], strides[1], strides[2]],
1045            },
1046            4 => NdDims::D4 {
1047                shape: [shape[0], shape[1], shape[2], shape[3]],
1048                strides: [strides[0], strides[1], strides[2], strides[3]],
1049            },
1050            5 => NdDims::D5 {
1051                shape: [shape[0], shape[1], shape[2], shape[3], shape[4]],
1052                strides: [strides[0], strides[1], strides[2], strides[3], strides[4]],
1053            },
1054            _ => NdDims::Dn {
1055                shape: shape.into(),
1056                strides: strides.into(),
1057            },
1058        }
1059    }
1060
1061    /// Number of dimensions.
1062    #[inline]
1063    pub(crate) fn ndim(&self) -> usize {
1064        match self {
1065            NdDims::D1 { .. } => 1,
1066            NdDims::D2 { .. } => 2,
1067            NdDims::D3 { .. } => 3,
1068            NdDims::D4 { .. } => 4,
1069            NdDims::D5 { .. } => 5,
1070            NdDims::Dn { shape, .. } => shape.len(),
1071        }
1072    }
1073
1074    /// Shape as a slice.
1075    #[inline]
1076    pub(crate) fn shape(&self) -> &[usize] {
1077        match self {
1078            NdDims::D1 { shape, .. } => shape,
1079            NdDims::D2 { shape, .. } => shape,
1080            NdDims::D3 { shape, .. } => shape,
1081            NdDims::D4 { shape, .. } => shape,
1082            NdDims::D5 { shape, .. } => shape,
1083            NdDims::Dn { shape, .. } => shape,
1084        }
1085    }
1086
1087    /// Strides as a slice.
1088    #[inline]
1089    pub(crate) fn strides(&self) -> &[usize] {
1090        match self {
1091            NdDims::D1 { strides, .. } => strides,
1092            NdDims::D2 { strides, .. } => strides,
1093            NdDims::D3 { strides, .. } => strides,
1094            NdDims::D4 { strides, .. } => strides,
1095            NdDims::D5 { strides, .. } => strides,
1096            NdDims::Dn { strides, .. } => strides,
1097        }
1098    }
1099
1100    /// Total logical element count i.e. the product of all dimensions.
1101    #[inline]
1102    pub(crate) fn len(&self) -> usize {
1103        self.shape().iter().product()
1104    }
1105}
1106
1107// ****************************************************************
1108// Axis selection input
1109// ****************************************************************
1110
1111/// Build an axis-selection slice from mixed indices and ranges.
1112///
1113/// Each entry is any [`DataSelector`](crate::traits::selection::DataSelector) -
1114/// a single index collapses the dimension, and a contiguous range keeps it.
1115///
1116/// # Example
1117/// ```ignore
1118/// arr.slice(nd![0..3, 2, 1..4])
1119/// ```
1120#[macro_export]
1121macro_rules! nd {
1122    ($($sel:expr),+ $(,)?) => {
1123        &[$(&$sel as &dyn $crate::traits::selection::DataSelector),+]
1124    };
1125}
1126
1127// ****************************************************************
1128// Stride computation
1129// ****************************************************************
1130
1131/// Compute compact column-major strides with no inter-dimension padding.
1132///
1133/// For a shape `[a, b, c]`, the strides are `[1, a, a * b]`. The buffer is
1134/// fully contiguous, so DLPack consumers receive a dense tensor and range
1135/// indexing over the outermost axis reads logical data with no gaps. The
1136/// backing allocation start remains 64-byte aligned through `Vec64`.
1137pub(crate) fn col_major_strides(shape: &[usize]) -> Vec<usize> {
1138    let mut strides = Vec::with_capacity(shape.len());
1139    if shape.is_empty() {
1140        return strides;
1141    }
1142    strides.push(1);
1143    for d in 1..shape.len() {
1144        strides.push(strides[d - 1] * shape[d - 1]);
1145    }
1146    strides
1147}
1148
1149/// Total buffer length required for a given shape and strides.
1150pub(crate) fn buffer_len(shape: &[usize], strides: &[usize]) -> usize {
1151    if shape.iter().any(|&d| d == 0) {
1152        return 0;
1153    }
1154    // The last element is at sum((shape[d]-1) * strides[d]) for all d.
1155    // Buffer must hold one past that.
1156    let max_offset: usize = shape.iter()
1157        .zip(strides.iter())
1158        .map(|(&s, &st)| (s - 1) * st)
1159        .sum();
1160    max_offset + 1
1161}
1162
1163// ****************************************************************
1164// IntoIterator
1165// ****************************************************************
1166
1167/// Iterating an NdArray yields `T` values in column-major order,
1168/// walking contiguous runs along axis 0 (the innermost dimension)
1169/// and advancing through higher dimensions. Each column/slice is
1170/// a sequential cache-friendly read with no per-element arithmetic.
1171impl<'a, T: Float> IntoIterator for &'a NdArray<T> {
1172    type Item = T;
1173    type IntoIter = NdArrayIter<'a, T>;
1174
1175    #[inline]
1176    fn into_iter(self) -> NdArrayIter<'a, T> {
1177        let shape = self.dims.shape();
1178        let strides = self.dims.strides();
1179        if shape.is_empty() {
1180            return NdArrayIter {
1181                buf: self.data.as_slice(),
1182                n_inner: 1,
1183                inner_stride: 1,
1184                run_offsets: vec![0],
1185                run_idx: 0,
1186                inner_idx: 0,
1187                total: 1,
1188                yielded: 0,
1189            };
1190        }
1191        let n_inner = shape[0];
1192
1193        // Number of contiguous runs = product of all dims except axis 0
1194        let n_runs: usize = shape[1..].iter().product();
1195
1196        // Build the starting offset of each contiguous run.
1197        // For 1D there is one run at offset 0.
1198        // For 2D these are just [0, stride1, 2*stride1, ...].
1199        // For N-D we walk the outer indices in column-major order.
1200        let mut run_offsets = Vec::with_capacity(n_runs);
1201        if shape.len() <= 1 {
1202            run_offsets.push(0);
1203        } else {
1204            let outer_shape = &shape[1..];
1205            let outer_strides = &strides[1..];
1206            let mut outer_indices = vec![0usize; outer_shape.len()];
1207            for _ in 0..n_runs {
1208                let off: usize = outer_indices.iter()
1209                    .zip(outer_strides.iter())
1210                    .map(|(&i, &s)| i * s)
1211                    .sum();
1212                run_offsets.push(off);
1213                // Advance outer indices (column-major)
1214                let mut carry = true;
1215                for d in 0..outer_shape.len() {
1216                    if carry {
1217                        outer_indices[d] += 1;
1218                        if outer_indices[d] < outer_shape[d] {
1219                            carry = false;
1220                        } else {
1221                            outer_indices[d] = 0;
1222                        }
1223                    }
1224                }
1225            }
1226        }
1227
1228        NdArrayIter {
1229            buf: self.data.as_slice(),
1230            n_inner,
1231            inner_stride: strides[0],
1232            run_offsets,
1233            run_idx: 0,
1234            inner_idx: 0,
1235            total: self.len(),
1236            yielded: 0,
1237        }
1238    }
1239}
1240
1241/// Consuming iterator - collects logical elements then iterates.
1242impl<T: Float> IntoIterator for NdArray<T> {
1243    type Item = T;
1244    type IntoIter = std::vec::IntoIter<T>;
1245
1246    #[inline]
1247    fn into_iter(self) -> std::vec::IntoIter<T> {
1248        let v: Vec<T> = (&self).into_iter().collect();
1249        v.into_iter()
1250    }
1251}
1252
1253/// Iterator over NdArray elements in column-major order.
1254///
1255/// When `inner_stride` is 1 (the normal case), walks contiguous runs
1256/// along axis 0 with sequential memory reads. When `inner_stride` > 1
1257/// (e.g. after collapsing axis 0 via slicing), steps through strided
1258/// elements within each run.
1259pub struct NdArrayIter<'a, T> {
1260    pub(crate) buf: &'a [T],
1261    pub(crate) n_inner: usize,
1262    pub(crate) inner_stride: usize,
1263    pub(crate) run_offsets: Vec<usize>,
1264    pub(crate) run_idx: usize,
1265    pub(crate) inner_idx: usize,
1266    pub(crate) total: usize,
1267    pub(crate) yielded: usize,
1268}
1269
1270impl<'a, T: Float> Iterator for NdArrayIter<'a, T> {
1271    type Item = T;
1272
1273    #[inline]
1274    fn next(&mut self) -> Option<T> {
1275        if self.yielded >= self.total { return None; }
1276
1277        let val = self.buf[self.run_offsets[self.run_idx] + self.inner_idx * self.inner_stride];
1278        self.yielded += 1;
1279        self.inner_idx += 1;
1280        if self.inner_idx >= self.n_inner {
1281            self.inner_idx = 0;
1282            self.run_idx += 1;
1283        }
1284        Some(val)
1285    }
1286
1287    #[inline]
1288    fn size_hint(&self) -> (usize, Option<usize>) {
1289        let r = self.total - self.yielded;
1290        (r, Some(r))
1291    }
1292}
1293
1294impl<'a, T: Float> ExactSizeIterator for NdArrayIter<'a, T> {}
1295
1296// ****************************************************************
1297// Internal helpers
1298// ****************************************************************
1299
1300/// Compute flat buffer offset for an N-dimensional index.
1301/// Panics on rank mismatch or an out-of-bounds index, in release
1302/// builds included, since a wrong offset can silently read or write
1303/// another element's slot.
1304#[inline]
1305pub(crate) fn offset_of_impl(indices: &[usize], shape: &[usize], strides: &[usize]) -> usize {
1306    assert_eq!(
1307        indices.len(),
1308        shape.len(),
1309        "NdArray: {} indices for a {}D array",
1310        indices.len(),
1311        shape.len()
1312    );
1313    let mut offset = 0;
1314    for d in 0..shape.len() {
1315        assert!(
1316            indices[d] < shape[d],
1317            "NdArray: index {} out of bounds for dim {} (size {})",
1318            indices[d], d, shape[d]
1319        );
1320        offset += indices[d] * strides[d];
1321    }
1322    offset
1323}
1324
1325// ****************************************************************
1326// Trait implementations
1327// ****************************************************************
1328
1329impl<T: Float> Shape for NdArray<T> {
1330    fn shape(&self) -> ShapeDim {
1331        match self.dims.ndim() {
1332            0 => ShapeDim::Rank0(1),
1333            1 => ShapeDim::Rank1(self.dims.shape()[0]),
1334            2 => ShapeDim::Rank2 {
1335                rows: self.dims.shape()[0],
1336                cols: self.dims.shape()[1],
1337            },
1338            _ => ShapeDim::RankN(self.dims.shape().to_vec()),
1339        }
1340    }
1341}
1342
1343impl<T: Float> Concatenate for NdArray<T> {
1344    /// Concatenate along axis 0. All other dimensions must match.
1345    fn concat(self, other: Self) -> Result<Self, MinarrowError> {
1346        let s1 = self.dims.shape();
1347        let s2 = other.dims.shape();
1348        if s1.len() != s2.len() {
1349            return Err(MinarrowError::IncompatibleTypeError {
1350                from: "NdArray",
1351                to: "NdArray",
1352                message: Some(format!(
1353                    "Cannot concatenate {}D and {}D arrays", s1.len(), s2.len()
1354                )),
1355            });
1356        }
1357        if s1.is_empty() {
1358            return Err(MinarrowError::ShapeError {
1359                message: "Cannot concatenate rank-zero arrays along axis 0".to_string(),
1360            });
1361        }
1362        for d in 1..s1.len() {
1363            if s1[d] != s2[d] {
1364                return Err(MinarrowError::IncompatibleTypeError {
1365                    from: "NdArray",
1366                    to: "NdArray",
1367                    message: Some(format!(
1368                        "Dimension {} mismatch: {} vs {}", d, s1[d], s2[d]
1369                    )),
1370                });
1371            }
1372        }
1373
1374        let mut new_shape: Vec<usize> = s1.to_vec();
1375        new_shape[0] += s2[0];
1376
1377        // The result name joins both sides, matching Table's concat.
1378        let name = match (&self.name, &other.name) {
1379            (Some(a), Some(b)) => Some(format!("{}+{}", a, b)),
1380            (Some(a), None) => Some(a.clone()),
1381            (None, Some(b)) => Some(b.clone()),
1382            (None, None) => None,
1383        };
1384
1385        // 1D: plain append
1386        if s1.len() == 1 {
1387            let mut flat = Vec64::with_capacity(self.len() + other.len());
1388            flat.extend(&self);
1389            flat.extend(&other);
1390            let mut result = NdArray::from_slice(&flat, &new_shape);
1391            result.name = name;
1392            return Ok(result);
1393        }
1394
1395        // Fast path for contiguous 2D: interleave columns with memcpy
1396        if s1.len() == 2 && self.dims.strides()[0] == 1 && other.dims.strides()[0] == 1 {
1397            let new_dims = NdDims::from_shape(&new_shape);
1398            let new_stride = new_dims.strides()[1];
1399            let total = buffer_len(&new_shape, new_dims.strides());
1400            let mut buf = Vec64::with_capacity(total);
1401            buf.0.resize(total, T::default());
1402            let dst = buf.as_mut_slice();
1403            for c in 0..s1[1] {
1404                let dst_start = c * new_stride;
1405                dst[dst_start..dst_start + s1[0]].copy_from_slice(self.col(c));
1406                dst[dst_start + s1[0]..dst_start + s1[0] + s2[0]].copy_from_slice(other.col(c));
1407            }
1408            return Ok(NdArray {
1409                data: Arc::new(Buffer::from_vec64(buf)),
1410                dims: new_dims,
1411                name,
1412            });
1413        }
1414
1415        // General case: interleave the axis-0 run of each operand per
1416        // outer index, walking both sources in column-major logical order.
1417        let n1 = s1[0];
1418        let n2 = s2[0];
1419        let run_len = n1 + n2;
1420        let n_runs: usize = s1[1..].iter().product();
1421        let new_dims = NdDims::from_shape(&new_shape);
1422        let total = buffer_len(&new_shape, new_dims.strides());
1423        let mut buf = Vec64::with_capacity(total);
1424        buf.0.resize(total, T::default());
1425        let dst = buf.as_mut_slice();
1426        let mut it_a = (&self).into_iter();
1427        let mut it_b = (&other).into_iter();
1428        for r in 0..n_runs {
1429            let base = r * run_len;
1430            for i in 0..n1 {
1431                dst[base + i] = it_a.next().unwrap();
1432            }
1433            for i in 0..n2 {
1434                dst[base + n1 + i] = it_b.next().unwrap();
1435            }
1436        }
1437        Ok(NdArray {
1438            data: Arc::new(Buffer::from_vec64(buf)),
1439            dims: new_dims,
1440            name,
1441        })
1442    }
1443}
1444
1445// *** Axis selection: arr.s(nd![1..4, 2]) *************************
1446
1447/// Selection across every axis at once, delegating to `slice`. Single
1448/// indices collapse their dimension, and contiguous ranges keep it.
1449/// Zero-copy.
1450#[cfg(all(feature = "views", feature = "select"))]
1451impl<T: Float> AxisSelection for NdArray<T> {
1452    type View = NdArrayV<T>;
1453
1454    fn s(&self, selection: &[&dyn DataSelector]) -> NdArrayV<T> {
1455        self.slice(selection)
1456    }
1457
1458    fn get_axis_count(&self) -> usize {
1459        self.ndim()
1460    }
1461}
1462
1463// *** Row selection: arr.r(0..10) *********************************
1464
1465/// Axis-0 observation selection. Contiguous ranges return a zero-copy
1466/// window view. Index arrays gather the selected observations into an
1467/// owned array wrapped in a full view, matching `Table`'s behaviour.
1468#[cfg(all(feature = "views", feature = "select"))]
1469impl<T: Float> RowSelection for NdArray<T> {
1470    type View = NdArrayV<T>;
1471
1472    fn r<S: DataSelector>(&self, selection: S) -> NdArrayV<T> {
1473        assert!(self.ndim() > 0, "row selection requires an axis 0");
1474        let n_obs = self.shape()[0];
1475        let indices = selection.resolve_indices(n_obs);
1476        if selection.is_contiguous() {
1477            let start = indices.first().copied().unwrap_or(0);
1478            let ranges: Vec<Range<usize>> = std::iter::once(start..start + indices.len())
1479                .chain(self.shape()[1..].iter().map(|&n| 0..n))
1480                .collect();
1481            let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect();
1482            return self.slice(&refs);
1483        }
1484        NdArrayV::from_ndarray(gather_obs_impl(
1485            &indices,
1486            self.shape(),
1487            self.name.clone(),
1488            |idx| self.get(idx),
1489        ))
1490    }
1491
1492    fn get_row_count(&self) -> usize {
1493        self.n_obs()
1494    }
1495}
1496
1497/// Materialise selected axis-0 observations into a compact owned array.
1498/// Walks the output positions in column-major order, reading each source
1499/// element through the provided accessor, so any stride layout gathers
1500/// correctly.
1501#[cfg(all(feature = "views", feature = "select"))]
1502pub(crate) fn gather_obs_impl<T: Float>(
1503    indices: &[usize],
1504    shape: &[usize],
1505    name: Option<String>,
1506    get: impl Fn(&[usize]) -> T,
1507) -> NdArray<T> {
1508    let mut out_shape = shape.to_vec();
1509    out_shape[0] = indices.len();
1510    let total: usize = out_shape.iter().product();
1511
1512    let mut flat = Vec64::with_capacity(total);
1513    let ndim = shape.len();
1514    let mut idx = vec![0usize; ndim];
1515    let inner_runs: usize = shape[1..].iter().product::<usize>().max(1);
1516    for _ in 0..inner_runs {
1517        for &obs in indices {
1518            idx[0] = obs;
1519            flat.push(get(&idx));
1520        }
1521        // Advance the inner multi-index in column-major order.
1522        let mut carry = true;
1523        for d in 1..ndim {
1524            if carry {
1525                idx[d] += 1;
1526                if idx[d] < shape[d] {
1527                    carry = false;
1528                } else {
1529                    idx[d] = 0;
1530                }
1531            }
1532        }
1533    }
1534    let mut result = NdArray::from_slice(&flat, &out_shape);
1535    result.name = name;
1536    result
1537}
1538
1539// *** Bracket indexing: arr[col][row] ******************************
1540
1541/// `arr[i]` selects along the outermost stored axis.
1542///
1543/// For 1D, returns a single-element slice.
1544/// For 2D (column-major), `arr[col]` returns the contiguous column
1545/// as `&[f64]`, so `arr[col][row]` gives `&f64`.
1546impl<T: Float> Index<usize> for NdArray<T> {
1547    type Output = [T];
1548
1549    #[inline]
1550    fn index(&self, idx: usize) -> &[T] {
1551        let shape = self.dims.shape();
1552        let strides = self.dims.strides();
1553        match shape.len() {
1554            0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1555            1 => {
1556                assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
1557                &self.data.as_slice()[idx..idx + 1]
1558            }
1559            2 => {
1560                assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
1561                let start = idx * strides[1];
1562                &self.data.as_slice()[start..start + shape[0]]
1563            }
1564            n => {
1565                // Index the outermost axis (last), return the contiguous inner slab
1566                assert!(
1567                    self.is_contiguous(),
1568                    "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
1569                );
1570                let last = n - 1;
1571                assert!(idx < shape[last], "index out of bounds for axis {}", last);
1572                let start = idx * strides[last];
1573                &self.data.as_slice()[start..start + strides[last]]
1574            }
1575        }
1576    }
1577}
1578
1579impl<T: Float> IndexMut<usize> for NdArray<T> {
1580    #[inline]
1581    fn index_mut(&mut self, idx: usize) -> &mut [T] {
1582        let shape = self.dims.shape().to_vec();
1583        let strides = self.dims.strides().to_vec();
1584        match shape.len() {
1585            0 => panic!("NdArray: a rank-zero array has no axis to index with usize"),
1586            1 => {
1587                assert!(idx < shape[0], "NdArray: index {} out of bounds (size {})", idx, shape[0]);
1588                &mut Arc::make_mut(&mut self.data).as_mut_slice()[idx..idx + 1]
1589            }
1590            2 => {
1591                assert!(idx < shape[1], "NdArray: column {} out of bounds (n_cols {})", idx, shape[1]);
1592                let start = idx * strides[1];
1593                let n_rows = shape[0];
1594                &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + n_rows]
1595            }
1596            n => {
1597                assert!(
1598                    self.is_contiguous(),
1599                    "outermost-axis indexing on 3D+ requires a contiguous layout, use slice() for strided access"
1600                );
1601                let last = n - 1;
1602                assert!(idx < shape[last], "index out of bounds for axis {}", last);
1603                let start = idx * strides[last];
1604                &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..start + strides[last]]
1605            }
1606        }
1607    }
1608}
1609
1610// *** Range indexing: arr[1..4] ************************************
1611
1612/// `arr[start..end]` selects a contiguous range along the outermost axis.
1613///
1614/// For 1D, returns the element slice directly.
1615/// For 2D and above, selects a range of outermost-axis entries as one
1616/// contiguous slab of logical data. Requires a contiguous layout, since a
1617/// padded or transposed stride pattern has no gap-free slab to return.
1618/// Non-contiguous arrays panic with guidance to use `slice()`.
1619impl<T: Float> Index<Range<usize>> for NdArray<T> {
1620    type Output = [T];
1621
1622    #[inline]
1623    fn index(&self, range: Range<usize>) -> &[T] {
1624        let shape = self.dims.shape();
1625        let strides = self.dims.strides();
1626        match shape.len() {
1627            0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1628            1 => &self.data.as_slice()[range],
1629            _ => {
1630                assert!(
1631                    self.is_contiguous(),
1632                    "range indexing requires a contiguous layout, use slice() for strided access"
1633                );
1634                let last = shape.len() - 1;
1635                assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
1636                let start = range.start * strides[last];
1637                let end = range.end * strides[last];
1638                &self.data.as_slice()[start..end]
1639            }
1640        }
1641    }
1642}
1643
1644impl<T: Float> IndexMut<Range<usize>> for NdArray<T> {
1645    #[inline]
1646    fn index_mut(&mut self, range: Range<usize>) -> &mut [T] {
1647        let shape = self.dims.shape().to_vec();
1648        let strides = self.dims.strides().to_vec();
1649        match shape.len() {
1650            0 => panic!("NdArray: a rank-zero array has no axis to range-index"),
1651            1 => &mut Arc::make_mut(&mut self.data).as_mut_slice()[range],
1652            _ => {
1653                assert!(
1654                    self.is_contiguous(),
1655                    "range indexing requires a contiguous layout, use slice() for strided access"
1656                );
1657                let last = shape.len() - 1;
1658                assert!(range.end <= shape[last], "NdArray: range end {} out of bounds (size {})", range.end, shape[last]);
1659                let start = range.start * strides[last];
1660                let end = range.end * strides[last];
1661                &mut Arc::make_mut(&mut self.data).as_mut_slice()[start..end]
1662            }
1663        }
1664    }
1665}
1666
1667impl<T: Float> Index<RangeFrom<usize>> for NdArray<T> {
1668    type Output = [T];
1669
1670    #[inline]
1671    fn index(&self, range: RangeFrom<usize>) -> &[T] {
1672        assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
1673        let last_dim = self.dims.shape().len() - 1;
1674        let end = self.dims.shape()[last_dim];
1675        &self[range.start..end]
1676    }
1677}
1678
1679impl<T: Float> Index<RangeTo<usize>> for NdArray<T> {
1680    type Output = [T];
1681
1682    #[inline]
1683    fn index(&self, range: RangeTo<usize>) -> &[T] {
1684        &self[0..range.end]
1685    }
1686}
1687
1688impl<T: Float> Index<RangeFull> for NdArray<T> {
1689    type Output = [T];
1690
1691    #[inline]
1692    fn index(&self, _: RangeFull) -> &[T] {
1693        assert!(self.ndim() > 0, "NdArray: a rank-zero array has no axis to range-index");
1694        let last_dim = self.dims.shape().len() - 1;
1695        let end = self.dims.shape()[last_dim];
1696        &self[0..end]
1697    }
1698}
1699
1700// *** Tuple indexing **********************************************
1701
1702impl<T: Float> Index<()> for NdArray<T> {
1703    type Output = T;
1704    #[inline]
1705    fn index(&self, (): ()) -> &T {
1706        &self.data.as_slice()[self.offset_of(&[])]
1707    }
1708}
1709
1710impl<T: Float> Index<(usize,)> for NdArray<T> {
1711    type Output = T;
1712    #[inline]
1713    fn index(&self, (i,): (usize,)) -> &T {
1714        &self.data.as_slice()[self.offset_of(&[i])]
1715    }
1716}
1717
1718impl<T: Float> Index<(usize, usize)> for NdArray<T> {
1719    type Output = T;
1720    #[inline]
1721    fn index(&self, (i, j): (usize, usize)) -> &T {
1722        &self.data.as_slice()[self.offset_of(&[i, j])]
1723    }
1724}
1725
1726impl<T: Float> Index<(usize, usize, usize)> for NdArray<T> {
1727    type Output = T;
1728    #[inline]
1729    fn index(&self, (i, j, k): (usize, usize, usize)) -> &T {
1730        &self.data.as_slice()[self.offset_of(&[i, j, k])]
1731    }
1732}
1733
1734impl<T: Float> Index<(usize, usize, usize, usize)> for NdArray<T> {
1735    type Output = T;
1736    #[inline]
1737    fn index(&self, (i, j, k, l): (usize, usize, usize, usize)) -> &T {
1738        &self.data.as_slice()[self.offset_of(&[i, j, k, l])]
1739    }
1740}
1741
1742impl<T: Float> Index<(usize, usize, usize, usize, usize)> for NdArray<T> {
1743    type Output = T;
1744    #[inline]
1745    fn index(&self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &T {
1746        &self.data.as_slice()[self.offset_of(&[i, j, k, l, m])]
1747    }
1748}
1749
1750impl<T: Float> IndexMut<(usize,)> for NdArray<T> {
1751    #[inline]
1752    fn index_mut(&mut self, (i,): (usize,)) -> &mut T {
1753        let off = self.offset_of(&[i]);
1754        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1755    }
1756}
1757
1758impl<T: Float> IndexMut<()> for NdArray<T> {
1759    #[inline]
1760    fn index_mut(&mut self, (): ()) -> &mut T {
1761        let off = self.offset_of(&[]);
1762        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1763    }
1764}
1765
1766impl<T: Float> IndexMut<(usize, usize)> for NdArray<T> {
1767    #[inline]
1768    fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut T {
1769        let off = self.offset_of(&[i, j]);
1770        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1771    }
1772}
1773
1774impl<T: Float> IndexMut<(usize, usize, usize)> for NdArray<T> {
1775    #[inline]
1776    fn index_mut(&mut self, (i, j, k): (usize, usize, usize)) -> &mut T {
1777        let off = self.offset_of(&[i, j, k]);
1778        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1779    }
1780}
1781
1782impl<T: Float> IndexMut<(usize, usize, usize, usize)> for NdArray<T> {
1783    #[inline]
1784    fn index_mut(&mut self, (i, j, k, l): (usize, usize, usize, usize)) -> &mut T {
1785        let off = self.offset_of(&[i, j, k, l]);
1786        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1787    }
1788}
1789
1790impl<T: Float> IndexMut<(usize, usize, usize, usize, usize)> for NdArray<T> {
1791    #[inline]
1792    fn index_mut(&mut self, (i, j, k, l, m): (usize, usize, usize, usize, usize)) -> &mut T {
1793        let off = self.offset_of(&[i, j, k, l, m]);
1794        &mut Arc::make_mut(&mut self.data).as_mut_slice()[off]
1795    }
1796}
1797
1798// *** From conversions ********************************************
1799
1800/// 1D from a flat slice.
1801impl<T: Float> From<&[T]> for NdArray<T> {
1802    fn from(data: &[T]) -> Self {
1803        NdArray {
1804            data: Arc::new(Buffer::from_slice(data)),
1805            dims: NdDims::from_shape(&[data.len()]),
1806            name: None,
1807        }
1808    }
1809}
1810
1811/// 1D from owned Vec64.
1812impl<T: Float> From<Vec64<T>> for NdArray<T> {
1813    fn from(v: Vec64<T>) -> Self {
1814        let n = v.len();
1815        NdArray {
1816            data: Arc::new(Buffer::from_vec64(v)),
1817            dims: NdDims::from_shape(&[n]),
1818            name: None,
1819        }
1820    }
1821}
1822
1823/// 2D from column vectors.
1824impl<T: Float> From<&[Vec<T>]> for NdArray<T> {
1825    fn from(columns: &[Vec<T>]) -> Self {
1826        let n_cols = columns.len();
1827        if n_cols == 0 {
1828            return NdArray::new(&[0, 0]);
1829        }
1830        let n_rows = columns[0].len();
1831        for col in columns {
1832            assert_eq!(col.len(), n_rows, "Column length mismatch");
1833        }
1834        let shape = [n_rows, n_cols];
1835        let dims = NdDims::from_shape(&shape);
1836        let stride = dims.strides()[1];
1837        let total = buffer_len(&shape, dims.strides());
1838        let mut buf = Vec64::with_capacity(total);
1839        buf.0.resize(total, T::default());
1840        for (c, col) in columns.iter().enumerate() {
1841            let start = c * stride;
1842            buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col);
1843        }
1844        NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
1845    }
1846}
1847
1848/// 2D from FloatArray columns.
1849impl<T: Float> From<&[FloatArray<T>]> for NdArray<T> {
1850    fn from(columns: &[FloatArray<T>]) -> Self {
1851        let n_cols = columns.len();
1852        if n_cols == 0 {
1853            return NdArray::new(&[0, 0]);
1854        }
1855        let n_rows = columns[0].data.len();
1856        for col in columns {
1857            assert_eq!(col.data.len(), n_rows, "Column length mismatch");
1858        }
1859        let shape = [n_rows, n_cols];
1860        let dims = NdDims::from_shape(&shape);
1861        let stride = dims.strides()[1];
1862        let total = buffer_len(&shape, dims.strides());
1863        let mut buf = Vec64::with_capacity(total);
1864        buf.0.resize(total, T::default());
1865        for (c, col) in columns.iter().enumerate() {
1866            let start = c * stride;
1867            buf.as_mut_slice()[start..start + n_rows].copy_from_slice(col.data.as_slice());
1868        }
1869        NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name: None }
1870    }
1871}
1872
1873/// From Matrix - zero-copy, moves the Buffer straight across. The Matrix's
1874/// padded column stride carries through, so the resulting array reports
1875/// non-contiguous. Call `to_contiguous` to re-lay out compactly.
1876#[cfg(feature = "matrix")]
1877impl From<Matrix> for NdArray<f64> {
1878    fn from(mat: Matrix) -> Self {
1879        let shape = [mat.n_rows, mat.n_cols];
1880        let strides = [1, mat.stride];
1881        NdArray {
1882            data: Arc::new(mat.data),
1883            dims: NdDims::from_shape_and_strides(&shape, &strides),
1884            name: mat.name,
1885        }
1886    }
1887}
1888
1889/// TryFrom Table - extracts numeric columns, converts nulls to NaN.
1890impl TryFrom<&Table> for NdArray<f64> {
1891    type Error = MinarrowError;
1892
1893    fn try_from(table: &Table) -> Result<Self, Self::Error> {
1894        let n_cols = table.n_cols();
1895        let n_rows = table.n_rows;
1896        if n_cols == 0 {
1897            return Ok(NdArray::new(&[0, 0]));
1898        }
1899
1900        let shape = [n_rows, n_cols];
1901        let dims = NdDims::from_shape(&shape);
1902        let stride = dims.strides()[1];
1903        let total = buffer_len(&shape, dims.strides());
1904        let mut buf = Vec64::with_capacity(total);
1905        buf.0.resize(total, 0.0);
1906
1907        for (col_idx, fa) in table.cols.iter().enumerate() {
1908            let numeric = fa.array.try_num().map_err(|_| MinarrowError::TypeError {
1909                from: "non-numeric",
1910                to: "Float64",
1911                message: Some(format!("column {} is not numeric", col_idx)),
1912            })?;
1913            let f64_arr = numeric.try_f64()?;
1914            if f64_arr.data.len() != n_rows {
1915                return Err(MinarrowError::ColumnLengthMismatch {
1916                    col: col_idx,
1917                    expected: n_rows,
1918                    found: f64_arr.data.len(),
1919                });
1920            }
1921
1922            let start = col_idx * stride;
1923            let src = f64_arr.data.as_slice();
1924            let dst = &mut buf.as_mut_slice()[start..start + n_rows];
1925
1926            // Copy data, converting nulls to NaN
1927            match f64_arr.null_mask.as_ref() {
1928                Some(mask) => {
1929                    for i in 0..n_rows {
1930                        dst[i] = if mask.get(i) { src[i] } else { f64::NAN };
1931                    }
1932                }
1933                None => dst.copy_from_slice(src),
1934            }
1935        }
1936
1937        let name = if table.name.is_empty() { None } else { Some(table.name.clone()) };
1938        Ok(NdArray { data: Arc::new(Buffer::from_vec64(buf)), dims, name })
1939    }
1940}
1941
1942impl TryFrom<Table> for NdArray<f64> {
1943    type Error = MinarrowError;
1944    fn try_from(table: Table) -> Result<Self, Self::Error> {
1945        NdArray::try_from(&table)
1946    }
1947}
1948
1949// *** Debug *******************************************************
1950
1951impl<T: Float> fmt::Debug for NdArray<T> {
1952    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1953        write!(
1954            f, "NdArray{}: {:?} [{}D, col-major]",
1955            self.name.as_deref().map_or(String::new(), |n| format!(" '{}'", n)),
1956            self.dims.shape(),
1957            self.ndim(),
1958        )?;
1959        if self.ndim() == 0 {
1960            write!(f, "\n{:8.4}", self.get(&[]).to_f64().unwrap_or(f64::NAN))?;
1961        } else if self.ndim() == 2 {
1962            let shape = self.dims.shape();
1963            let max_rows = shape[0].min(6);
1964            let max_cols = shape[1].min(8);
1965            for r in 0..max_rows {
1966                write!(f, "\n[")?;
1967                for c in 0..max_cols {
1968                    write!(f, " {:8.4}", self.get(&[r, c]).to_f64().unwrap_or(f64::NAN))?;
1969                    if c < max_cols - 1 { write!(f, ",")?; }
1970                }
1971                if shape[1] > 8 { write!(f, " ...")?; }
1972                write!(f, " ]")?;
1973            }
1974            if shape[0] > 6 { write!(f, "\n...")?; }
1975        } else if self.ndim() == 1 {
1976            let n = self.dims.shape()[0].min(10);
1977            write!(f, "\n[")?;
1978            for i in 0..n {
1979                write!(f, " {:8.4}", self.get(&[i]).to_f64().unwrap_or(f64::NAN))?;
1980                if i < n - 1 { write!(f, ",")?; }
1981            }
1982            if self.dims.shape()[0] > 10 { write!(f, " ...")?; }
1983            write!(f, " ]")?;
1984        }
1985        Ok(())
1986    }
1987}
1988
1989#[cfg(test)]
1990mod tests {
1991    use super::*;
1992    use crate::StringArray;
1993    use crate::structs::bitmask::Bitmask;
1994
1995    // *** Row selection and apply *************************************
1996
1997    #[cfg(all(feature = "views", feature = "select"))]
1998    #[test]
1999    fn axis_selection_rank10() {
2000        let shape = [3, 4, 2, 5, 3, 2, 4, 3, 2, 5];
2001        let len: usize = shape.iter().product();
2002        let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
2003        let a = NdArray::from_slice(&data, &shape);
2004
2005        // Mixed ranges, single indices, and a full range across all ten axes.
2006        let v = a.s(nd![1..3, 0..2, 1, 2..5, .., 0..1, 1..3, 2, 0..2, 3..5]);
2007        assert_eq!(v.shape(), &[2, 2, 3, 3, 1, 2, 2, 2]);
2008        assert_eq!(
2009            v.get(&[0, 0, 0, 0, 0, 0, 0, 0]),
2010            a.get(&[1, 0, 1, 2, 0, 0, 1, 2, 0, 3])
2011        );
2012        assert_eq!(
2013            v.get(&[1, 1, 2, 2, 0, 1, 1, 1]),
2014            a.get(&[2, 1, 1, 4, 2, 0, 2, 2, 1, 4])
2015        );
2016    }
2017
2018    #[cfg(all(feature = "views", feature = "select"))]
2019    #[test]
2020    fn axis_selection_runtime_rank() {
2021        // Selections built at runtime for ranks beyond literal syntax.
2022        let mut shape = vec![1usize; 100];
2023        shape[0] = 3;
2024        shape[10] = 4;
2025        shape[50] = 5;
2026        shape[99] = 2;
2027        let len: usize = shape.iter().product();
2028        let data: Vec<f64> = (0..len).map(|i| i as f64).collect();
2029        let a = NdArray::from_slice(&data, &shape);
2030
2031        // Full range on every axis, then narrow three and collapse one.
2032        let mut sels: Vec<Box<dyn DataSelector>> =
2033            shape.iter().map(|&n| Box::new(0..n) as Box<dyn DataSelector>).collect();
2034        sels[0] = Box::new(1..3);
2035        sels[10] = Box::new(2usize);
2036        sels[50] = Box::new(1..4);
2037        let refs: Vec<&dyn DataSelector> = sels.iter().map(|s| s.as_ref()).collect();
2038
2039        let v = a.s(&refs);
2040        assert_eq!(v.ndim(), 99);
2041        assert_eq!(v.shape()[0], 2);
2042        assert_eq!(v.shape()[49], 3);
2043        assert_eq!(v.shape()[98], 2);
2044
2045        let mut view_idx = vec![0usize; 99];
2046        view_idx[0] = 1;
2047        view_idx[49] = 2;
2048        view_idx[98] = 1;
2049        let mut source_idx = vec![0usize; 100];
2050        source_idx[0] = 2;
2051        source_idx[10] = 2;
2052        source_idx[50] = 3;
2053        source_idx[99] = 1;
2054        assert_eq!(v.get(&view_idx), a.get(&source_idx));
2055    }
2056
2057    #[cfg(all(feature = "views", feature = "select"))]
2058    #[test]
2059    fn axis_selection_trait() {
2060        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2061        // Range keeps the axis, index collapses it.
2062        let v = a.s(nd![1..3, 1]);
2063        assert_eq!(v.shape(), &[2]);
2064        assert_eq!(v.get(&[0]), 5.0);
2065        assert_eq!(v.get(&[1]), 6.0);
2066        // Selection composes on the view.
2067        let sub = a.s(nd![0..3, 0..2]).s(nd![2, 0..2]);
2068        assert_eq!(sub.shape(), &[2]);
2069        assert_eq!(sub.get(&[0]), 3.0);
2070        assert_eq!(sub.get(&[1]), 6.0);
2071        assert_eq!(a.get_axis_count(), 2);
2072        assert_eq!(sub.get_axis_count(), 1);
2073    }
2074
2075    #[cfg(all(feature = "views", feature = "select"))]
2076    #[test]
2077    fn row_selection_contiguous() {
2078        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2079        let v = a.r(1..3);
2080        assert_eq!(v.shape(), &[2, 2]);
2081        assert_eq!(v.get(&[0, 0]), 2.0);
2082        assert_eq!(v.get(&[1, 1]), 6.0);
2083        // The row alias selects a single observation.
2084        let single = a.row(2);
2085        assert_eq!(single.shape(), &[1, 2]);
2086        assert_eq!(single.get(&[0, 1]), 6.0);
2087    }
2088
2089    #[cfg(all(feature = "views", feature = "select"))]
2090    #[test]
2091    fn row_selection_gathers_indices() {
2092        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2093        let v = a.r(&[2, 0]);
2094        assert_eq!(v.shape(), &[2, 2]);
2095        // Gathered rows follow selection order.
2096        assert_eq!(v.get(&[0, 0]), 3.0);
2097        assert_eq!(v.get(&[0, 1]), 6.0);
2098        assert_eq!(v.get(&[1, 0]), 1.0);
2099        assert_eq!(v.get(&[1, 1]), 4.0);
2100    }
2101
2102    #[test]
2103    fn apply_maps_elements() {
2104        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2105        a.set_name("m");
2106        let b = a.apply(|x| x * 10.0);
2107        assert_eq!(b.get(&[1, 1]), 40.0);
2108        assert_eq!(b.name.as_deref(), Some("m"));
2109        // The source is untouched.
2110        assert_eq!(a.get(&[1, 1]), 4.0);
2111    }
2112
2113    #[test]
2114    fn apply_mut_in_place() {
2115        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2116        a.apply_mut(|x| x + 0.5);
2117        assert_eq!((&a).into_iter().collect::<Vec<f64>>(), vec![1.5, 2.5, 3.5]);
2118    }
2119
2120    #[cfg(feature = "matrix")]
2121    #[test]
2122    fn apply_mut_non_contiguous_touches_logical_only() {
2123        // A Matrix-imported array carries stride padding. The logical walk
2124        // mutates only real elements.
2125        let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None);
2126        let mut a = NdArray::from(mat);
2127        assert!(!a.is_contiguous());
2128        a.apply_mut(|x| x * 2.0);
2129        assert_eq!(a.get(&[0, 0]), 2.0);
2130        assert_eq!(a.get(&[2, 1]), 12.0);
2131    }
2132
2133    #[cfg(feature = "views")]
2134    #[test]
2135    fn apply_axis_collapses_axis() {
2136        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2137        // Sum each column lane (axis 0) - output shape [2].
2138        let col_sums = a.apply_axis(0, |lane| (&lane).into_iter().sum());
2139        assert_eq!(col_sums.shape(), &[2]);
2140        assert_eq!(col_sums.get(&[0]), 6.0);
2141        assert_eq!(col_sums.get(&[1]), 15.0);
2142        // Sum each row lane (axis 1) - output shape [3].
2143        let row_sums = a.apply_axis(1, |lane| (&lane).into_iter().sum());
2144        assert_eq!(row_sums.shape(), &[3]);
2145        assert_eq!(row_sums.get(&[0]), 5.0);
2146        assert_eq!(row_sums.get(&[2]), 9.0);
2147    }
2148
2149    #[cfg(feature = "views")]
2150    #[test]
2151    fn apply_axis_3d() {
2152        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2153        let a = NdArray::from_slice(&data, &[2, 3, 4]);
2154        let maxes = a.apply_axis(1, |lane| {
2155            (&lane).into_iter().fold(f64::MIN, f64::max)
2156        });
2157        assert_eq!(maxes.shape(), &[2, 4]);
2158        // Lane over axis 1 at [0, .., 0] holds values at [0,j,0].
2159        assert_eq!(maxes.get(&[0, 0]), a.get(&[0, 2, 0]));
2160        assert_eq!(maxes.get(&[1, 3]), a.get(&[1, 2, 3]));
2161    }
2162
2163    // ****************************************************************
2164    // Construction
2165    // ****************************************************************
2166
2167    #[test]
2168    fn new_zeroed_1d() {
2169        let a = NdArray::<f64>::new(&[5]);
2170        assert_eq!(a.ndim(), 1);
2171        assert_eq!(a.shape(), &[5]);
2172        assert_eq!(a.len(), 5);
2173        assert!(!a.is_empty());
2174        for v in &a { assert_eq!(v, 0.0); }
2175    }
2176
2177    #[test]
2178    fn new_zeroed_2d() {
2179        let a = NdArray::<f64>::new(&[3, 4]);
2180        assert_eq!(a.ndim(), 2);
2181        assert_eq!(a.shape(), &[3, 4]);
2182        assert_eq!(a.len(), 12);
2183        for v in &a { assert_eq!(v, 0.0); }
2184    }
2185
2186    #[test]
2187    fn new_zeroed_3d() {
2188        let a = NdArray::<f64>::new(&[2, 3, 4]);
2189        assert_eq!(a.ndim(), 3);
2190        assert_eq!(a.shape(), &[2, 3, 4]);
2191        assert_eq!(a.len(), 24);
2192    }
2193
2194    #[test]
2195    fn new_zeroed_5d() {
2196        let a = NdArray::<f64>::new(&[2, 3, 4, 5, 6]);
2197        assert_eq!(a.ndim(), 5);
2198        assert_eq!(a.len(), 720);
2199    }
2200
2201    #[test]
2202    fn new_zeroed_6d_heap() {
2203        let a = NdArray::<f64>::new(&[2, 3, 2, 2, 2, 2]);
2204        assert_eq!(a.ndim(), 6);
2205        assert_eq!(a.len(), 96);
2206    }
2207
2208    #[test]
2209    fn new_named() {
2210        let a = NdArray::<f64>::new_named(&[3, 3], "covariance");
2211        assert_eq!(a.name.as_deref(), Some("covariance"));
2212    }
2213
2214    #[test]
2215    fn from_slice_1d() {
2216        let data = [1.0, 2.0, 3.0, 4.0, 5.0];
2217        let a = NdArray::from_slice(&data, &[5]);
2218        assert_eq!(a.len(), 5);
2219        let vals: Vec<f64> = (&a).into_iter().collect();
2220        assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
2221    }
2222
2223    #[test]
2224    fn from_slice_2d_column_major() {
2225        let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2226        let a = NdArray::from_slice(&data, &[3, 2]);
2227        assert_eq!(a.get(&[0, 0]), 1.0);
2228        assert_eq!(a.get(&[1, 0]), 2.0);
2229        assert_eq!(a.get(&[2, 0]), 3.0);
2230        assert_eq!(a.get(&[0, 1]), 4.0);
2231        assert_eq!(a.get(&[1, 1]), 5.0);
2232        assert_eq!(a.get(&[2, 1]), 6.0);
2233    }
2234
2235    #[test]
2236    fn fill_and_ones() {
2237        let a = NdArray::fill(&[2, 3], 7.0);
2238        assert_eq!(a.len(), 6);
2239        for v in &a { assert_eq!(v, 7.0); }
2240
2241        let b = NdArray::<f64>::ones(&[4]);
2242        for v in &b { assert_eq!(v, 1.0); }
2243    }
2244
2245    #[test]
2246    fn eye_identity() {
2247        let a = NdArray::<f64>::eye(3);
2248        assert_eq!(a.shape(), &[3, 3]);
2249        assert_eq!(a.get(&[0, 0]), 1.0);
2250        assert_eq!(a.get(&[1, 1]), 1.0);
2251        assert_eq!(a.get(&[2, 2]), 1.0);
2252        assert_eq!(a.get(&[0, 1]), 0.0);
2253        assert_eq!(a.get(&[1, 0]), 0.0);
2254    }
2255
2256    #[test]
2257    fn linspace_basic() {
2258        let a = NdArray::<f64>::linspace(0.0, 1.0, 5);
2259        assert_eq!(a.shape(), &[5]);
2260        assert_eq!(a.get(&[0]), 0.0);
2261        assert_eq!(a.get(&[4]), 1.0);
2262        assert!((a.get(&[2]) - 0.5).abs() < 1e-15);
2263    }
2264
2265    #[test]
2266    fn arange_basic() {
2267        let a = NdArray::arange(0.0, 0.5, 4);
2268        assert_eq!(a.shape(), &[4]);
2269        assert_eq!(a.get(&[0]), 0.0);
2270        assert_eq!(a.get(&[1]), 0.5);
2271        assert_eq!(a.get(&[2]), 1.0);
2272        assert_eq!(a.get(&[3]), 1.5);
2273    }
2274
2275    #[test]
2276    fn from_vec64_moves_data() {
2277        let data: Vec64<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into();
2278        let a = NdArray::from_vec64(data, &[3, 2]);
2279        assert_eq!(a.shape(), &[3, 2]);
2280        assert_eq!(a.get(&[2, 1]), 6.0);
2281    }
2282
2283    #[cfg(feature = "views")]
2284    #[test]
2285    fn iter_obs_walks_observations() {
2286        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2287        let rows: Vec<(usize, Vec<f64>)> = a
2288            .iter_obs()
2289            .map(|(i, v)| (i, (&v).into_iter().collect()))
2290            .collect();
2291        assert_eq!(rows.len(), 3);
2292        assert_eq!(rows[0], (0, vec![1.0, 4.0]));
2293        assert_eq!(rows[2], (2, vec![3.0, 6.0]));
2294    }
2295
2296    #[test]
2297    fn equality_ignores_name() {
2298        let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2299        let mut b = a.clone();
2300        b.name = Some("named".to_string());
2301        assert_eq!(a, b);
2302    }
2303
2304    #[test]
2305    fn name_survives_reshape_and_concat() {
2306        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
2307        a.name = Some("px".to_string());
2308        assert_eq!(a.reshape(&[2, 2]).unwrap().name.as_deref(), Some("px"));
2309
2310        let mut b = NdArray::from_slice(&[5.0, 6.0], &[2]);
2311        b.name = Some("qty".to_string());
2312        let joined = a.concat(b).unwrap();
2313        assert_eq!(joined.name.as_deref(), Some("px+qty"));
2314    }
2315
2316    // ****************************************************************
2317    // Element access and indexing
2318    // ****************************************************************
2319
2320    #[test]
2321    fn get_set_1d() {
2322        let mut a = NdArray::new(&[3]);
2323        a.set(&[0], 10.0);
2324        a.set(&[1], 20.0);
2325        a.set(&[2], 30.0);
2326        assert_eq!(a.get(&[0]), 10.0);
2327        assert_eq!(a.get(&[1]), 20.0);
2328        assert_eq!(a.get(&[2]), 30.0);
2329    }
2330
2331    #[test]
2332    fn tuple_index_1d() {
2333        let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
2334        assert_eq!(a[(0,)], 10.0);
2335        assert_eq!(a[(1,)], 20.0);
2336        assert_eq!(a[(2,)], 30.0);
2337    }
2338
2339    #[test]
2340    fn tuple_index_2d() {
2341        let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2342        let a = NdArray::from_slice(&data, &[3, 2]);
2343        assert_eq!(a[(0, 0)], 1.0);
2344        assert_eq!(a[(2, 0)], 3.0);
2345        assert_eq!(a[(0, 1)], 4.0);
2346        assert_eq!(a[(2, 1)], 6.0);
2347    }
2348
2349    #[test]
2350    fn tuple_index_3d() {
2351        let data: Vec<f64> = (1..=12).map(|x| x as f64).collect();
2352        let a = NdArray::from_slice(&data, &[2, 3, 2]);
2353        assert_eq!(a[(0, 0, 0)], 1.0);
2354        assert_eq!(a[(1, 0, 0)], 2.0);
2355        assert_eq!(a[(0, 1, 0)], 3.0);
2356        assert_eq!(a[(1, 1, 0)], 4.0);
2357        assert_eq!(a[(0, 2, 0)], 5.0);
2358        assert_eq!(a[(1, 2, 0)], 6.0);
2359        assert_eq!(a[(0, 0, 1)], 7.0);
2360        assert_eq!(a[(1, 2, 1)], 12.0);
2361    }
2362
2363    #[test]
2364    fn index_mut_2d() {
2365        let mut a = NdArray::new(&[2, 2]);
2366        a[(0, 0)] = 1.0;
2367        a[(1, 0)] = 2.0;
2368        a[(0, 1)] = 3.0;
2369        a[(1, 1)] = 4.0;
2370        assert_eq!(a[(0, 0)], 1.0);
2371        assert_eq!(a[(1, 1)], 4.0);
2372    }
2373
2374    // ****************************************************************
2375    // Iteration
2376    // ****************************************************************
2377
2378    #[test]
2379    fn iter_1d() {
2380        let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
2381        let vals: Vec<f64> = (&a).into_iter().collect();
2382        assert_eq!(vals, vec![10.0, 20.0, 30.0]);
2383    }
2384
2385    #[test]
2386    fn iter_2d_column_major_order() {
2387        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2388        let vals: Vec<f64> = (&a).into_iter().collect();
2389        assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
2390    }
2391
2392    #[test]
2393    fn iter_2d_compact_layout() {
2394        // Compact strides place columns back to back, so iteration reads
2395        // the buffer straight through.
2396        let data: Vec<f64> = (1..=30).map(|x| x as f64).collect();
2397        let a = NdArray::from_slice(&data, &[10, 3]);
2398        assert_eq!(a.strides()[1], 10);
2399        let vals: Vec<f64> = (&a).into_iter().collect();
2400        assert_eq!(vals.len(), 30);
2401        assert_eq!(&vals[..10], &data[..10]);
2402        assert_eq!(&vals[10..20], &data[10..20]);
2403        assert_eq!(&vals[20..30], &data[20..30]);
2404    }
2405
2406    #[test]
2407    fn iter_3d_column_major_order() {
2408        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2409        let a = NdArray::from_slice(&data, &[2, 3, 4]);
2410        let vals: Vec<f64> = (&a).into_iter().collect();
2411        assert_eq!(vals.len(), 24);
2412        assert_eq!(vals, data);
2413    }
2414
2415    #[test]
2416    fn iter_exact_size() {
2417        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2418        let iter = (&a).into_iter();
2419        assert_eq!(iter.len(), 6);
2420    }
2421
2422    #[test]
2423    fn consuming_into_iter() {
2424        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2425        let vals: Vec<f64> = a.into_iter().collect();
2426        assert_eq!(vals, vec![1.0, 2.0, 3.0]);
2427    }
2428
2429    // ****************************************************************
2430    // Shape introspection
2431    // ****************************************************************
2432
2433    #[test]
2434    fn rank_zero_scalar_semantics() {
2435        let mut a = NdArray::from_slice(&[5.0], &[]);
2436        assert_eq!(a.ndim(), 0);
2437        assert!(a.shape().is_empty());
2438        assert!(a.strides().is_empty());
2439        assert_eq!(a.len(), 1);
2440        assert!(!a.is_empty());
2441        assert!(a.is_contiguous());
2442        assert_eq!(a[()], 5.0);
2443        assert_eq!((&a).into_iter().collect::<Vec<_>>(), vec![5.0]);
2444        assert_eq!(Shape::shape(&a), ShapeDim::Rank0(1));
2445
2446        a[()] = 7.0;
2447        assert_eq!(a[()], 7.0);
2448        assert!(a.transpose().shape().is_empty());
2449        assert_eq!(a.reshape(&[1]).unwrap().get(&[0]), 7.0);
2450    }
2451
2452    #[cfg(all(feature = "views", feature = "select"))]
2453    #[test]
2454    fn selecting_every_axis_can_produce_a_scalar_view() {
2455        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2456        let scalar = a.slice(&[&1usize, &1usize]);
2457        assert!(scalar.shape().is_empty());
2458        assert!(scalar.strides().is_empty());
2459        assert_eq!(scalar.len(), 1);
2460        assert_eq!(scalar.get(&[]), 4.0);
2461        assert_eq!((&scalar).into_iter().collect::<Vec<_>>(), vec![4.0]);
2462        assert_eq!(Shape::shape(&scalar), ShapeDim::Rank0(1));
2463    }
2464
2465    #[test]
2466    fn is_contiguous_default() {
2467        let a = NdArray::<f64>::new(&[3, 4]);
2468        assert!(a.is_contiguous());
2469    }
2470
2471    #[test]
2472    fn shape_trait_1d() {
2473        let a = NdArray::<f64>::new(&[5]);
2474        assert_eq!(Shape::shape(&a), ShapeDim::Rank1(5));
2475    }
2476
2477    #[test]
2478    fn shape_trait_2d() {
2479        let a = NdArray::<f64>::new(&[3, 4]);
2480        assert_eq!(Shape::shape(&a), ShapeDim::Rank2 { rows: 3, cols: 4 });
2481    }
2482
2483    #[test]
2484    fn shape_trait_3d() {
2485        let a = NdArray::<f64>::new(&[2, 3, 4]);
2486        assert_eq!(Shape::shape(&a), ShapeDim::RankN(vec![2, 3, 4]));
2487    }
2488
2489    // ****************************************************************
2490    // NaN handling
2491    // ****************************************************************
2492
2493    #[test]
2494    fn has_nan_false() {
2495        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2496        assert!(!a.has_nan());
2497        assert_eq!(a.nan_count(), 0);
2498    }
2499
2500    #[test]
2501    fn has_nan_true() {
2502        let a = NdArray::from_slice(&[1.0, f64::NAN, 3.0], &[3]);
2503        assert!(a.has_nan());
2504        assert_eq!(a.nan_count(), 1);
2505    }
2506
2507    #[test]
2508    fn has_nan_2d() {
2509        let a = NdArray::from_slice(&[1.0, 2.0, f64::NAN, 4.0, 5.0, f64::NAN], &[3, 2]);
2510        assert!(a.has_nan());
2511        assert_eq!(a.nan_count(), 2);
2512    }
2513
2514    // ****************************************************************
2515    // 2D axis access
2516    // ****************************************************************
2517
2518    #[test]
2519    fn col_access() {
2520        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2521        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2522        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2523    }
2524
2525    #[test]
2526    fn col_mut_access() {
2527        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2528        a.col_mut(0)[1] = 99.0;
2529        assert_eq!(a.col(0), &[1.0, 99.0, 3.0]);
2530    }
2531
2532    #[test]
2533    fn columns_access() {
2534        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2535        let cols = a.columns();
2536        assert_eq!(cols.len(), 2);
2537        assert_eq!(cols[0], &[1.0, 2.0, 3.0]);
2538        assert_eq!(cols[1], &[4.0, 5.0, 6.0]);
2539    }
2540
2541    #[test]
2542    fn columns_mut_access() {
2543        let mut a = NdArray::new(&[3, 2]);
2544        {
2545            let mut cols = a.columns_mut();
2546            cols[0].copy_from_slice(&[1.0, 2.0, 3.0]);
2547            cols[1].copy_from_slice(&[4.0, 5.0, 6.0]);
2548        }
2549        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2550        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2551    }
2552
2553    #[cfg(feature = "views")]
2554    #[test]
2555    fn obs_access() {
2556        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2557        let v = a.as_view();
2558        let obs0: Vec<f64> = (&v.obs(0)).into_iter().collect();
2559        let obs1: Vec<f64> = (&v.obs(1)).into_iter().collect();
2560        let obs2: Vec<f64> = (&v.obs(2)).into_iter().collect();
2561        assert_eq!(obs0, vec![1.0, 4.0]);
2562        assert_eq!(obs1, vec![2.0, 5.0]);
2563        assert_eq!(obs2, vec![3.0, 6.0]);
2564
2565        // Single-shot .obs() on NdArray also works
2566        assert_eq!((&a.obs(1)).into_iter().collect::<Vec<f64>>(), vec![2.0, 5.0]);
2567
2568        // 3D obs returns a 2D view
2569        let b = NdArray::from_slice(&(1..=24).map(|x| x as f64).collect::<Vec<_>>(), &[2, 3, 4]);
2570        let obs0_3d = b.obs(0);
2571        assert_eq!(obs0_3d.shape(), &[3, 4]);
2572    }
2573
2574    #[test]
2575    fn col_access_2d() {
2576        let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
2577        let a = NdArray::from_slice(&data, &[10, 2]);
2578        assert_eq!(a.col(0).len(), 10);
2579        assert_eq!(a.col(1).len(), 10);
2580        assert_eq!(a.col(0), &data[..10]);
2581        assert_eq!(a.col(1), &data[10..20]);
2582    }
2583
2584    // ****************************************************************
2585    // BLAS compatibility
2586    // ****************************************************************
2587
2588    #[test]
2589    fn blas_params() {
2590        let a = NdArray::<f64>::new(&[10, 5]);
2591        assert_eq!(a.m(), 10);
2592        assert_eq!(a.n(), 5);
2593        assert_eq!(a.lda(), 10);
2594    }
2595
2596    #[test]
2597    fn blas_params_aligned_rows() {
2598        let a = NdArray::<f64>::new(&[8, 3]);
2599        assert_eq!(a.m(), 8);
2600        assert_eq!(a.n(), 3);
2601        assert_eq!(a.lda(), 8);
2602    }
2603
2604    // ****************************************************************
2605    // Compact strides
2606    // ****************************************************************
2607
2608    #[test]
2609    fn compact_strides_2d() {
2610        for n_rows in 1..=20 {
2611            let a = NdArray::<f64>::new(&[n_rows, 3]);
2612            assert_eq!(a.strides()[1], n_rows);
2613            assert!(a.is_contiguous());
2614        }
2615    }
2616
2617    #[test]
2618    fn compact_strides_3d() {
2619        let a = NdArray::<f64>::new(&[10, 3, 4]);
2620        let strides = a.strides();
2621        assert_eq!(strides[0], 1);
2622        assert_eq!(strides[1], 10);
2623        assert_eq!(strides[2], 10 * 3);
2624        assert!(a.is_contiguous());
2625    }
2626
2627    // ****************************************************************
2628    // Reshape and transform
2629    // ****************************************************************
2630
2631    #[test]
2632    fn reshape_1d_to_2d() {
2633        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[6]);
2634        let b = a.reshape(&[3, 2]).unwrap();
2635        assert_eq!(b.shape(), &[3, 2]);
2636        assert_eq!(b.get(&[0, 0]), 1.0);
2637        assert_eq!(b.get(&[1, 0]), 2.0);
2638        assert_eq!(b.get(&[2, 0]), 3.0);
2639        assert_eq!(b.get(&[0, 1]), 4.0);
2640    }
2641
2642    #[test]
2643    fn reshape_size_mismatch() {
2644        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2645        assert!(a.reshape(&[2, 2]).is_err());
2646    }
2647
2648    #[test]
2649    fn transpose_2d() {
2650        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2651        let t = a.transpose();
2652        assert_eq!(t.shape(), &[2, 3]);
2653        assert_eq!(t.get(&[0, 0]), 1.0);
2654        assert_eq!(t.get(&[1, 0]), 4.0);
2655        assert_eq!(t.get(&[0, 1]), 2.0);
2656        assert_eq!(t.get(&[1, 1]), 5.0);
2657        assert_eq!(t.get(&[0, 2]), 3.0);
2658        assert_eq!(t.get(&[1, 2]), 6.0);
2659    }
2660
2661    #[test]
2662    fn transpose_3d_reverses_axes() {
2663        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
2664        let a = NdArray::from_slice(&data, &[2, 3, 4]);
2665        let t = a.transpose();
2666        assert_eq!(t.shape(), &[4, 3, 2]);
2667        assert!(t.is_contiguous());
2668        // Every element lands at its reversed index.
2669        for i in 0..2 {
2670            for j in 0..3 {
2671                for k in 0..4 {
2672                    assert_eq!(t.get(&[k, j, i]), a.get(&[i, j, k]));
2673                }
2674            }
2675        }
2676    }
2677
2678    #[test]
2679    fn transpose_1d_copies_through() {
2680        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2681        let t = a.transpose();
2682        assert_eq!(t.shape(), &[3]);
2683        assert_eq!(t, a);
2684    }
2685
2686    #[test]
2687    fn flatten() {
2688        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
2689        let flat = a.flatten();
2690        assert_eq!(flat.shape(), &[6]);
2691        let vals: Vec<f64> = (&flat).into_iter().collect();
2692        assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
2693    }
2694
2695    #[test]
2696    fn to_contiguous_noop() {
2697        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2698        let b = a.to_contiguous();
2699        assert_eq!(a, b);
2700    }
2701
2702    #[test]
2703    fn fill_with_contiguous() {
2704        let mut a = NdArray::new(&[3, 2]);
2705        a.fill_with(42.0);
2706        for v in &a { assert_eq!(v, 42.0); }
2707    }
2708
2709    // ****************************************************************
2710    // From conversions
2711    // ****************************************************************
2712
2713    #[test]
2714    fn from_f64_slice() {
2715        let a = NdArray::from(&[1.0, 2.0, 3.0][..]);
2716        assert_eq!(a.ndim(), 1);
2717        assert_eq!(a.len(), 3);
2718        assert_eq!(a[(0,)], 1.0);
2719    }
2720
2721    #[test]
2722    fn from_vec64() {
2723        let v: Vec64<f64> = vec![10.0, 20.0].into_iter().collect();
2724        let a = NdArray::from(v);
2725        assert_eq!(a.ndim(), 1);
2726        assert_eq!(a[(0,)], 10.0);
2727        assert_eq!(a[(1,)], 20.0);
2728    }
2729
2730    #[test]
2731    fn from_column_vecs() {
2732        let cols = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
2733        let a = NdArray::from(cols.as_slice());
2734        assert_eq!(a.shape(), &[3, 2]);
2735        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2736        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2737    }
2738
2739    #[test]
2740    fn from_float_arrays() {
2741        let c0 = FloatArray::from_slice(&[1.0, 2.0]);
2742        let c1 = FloatArray::from_slice(&[3.0, 4.0]);
2743        let a = NdArray::from([c0, c1].as_slice());
2744        assert_eq!(a.shape(), &[2, 2]);
2745        assert_eq!(a.col(0), &[1.0, 2.0]);
2746        assert_eq!(a.col(1), &[3.0, 4.0]);
2747    }
2748
2749    #[test]
2750    fn from_buffer_explicit_strides() {
2751        let mut buf = Vec64::with_capacity(16);
2752        buf.0.resize(16, 0.0);
2753        buf[0] = 1.0; buf[1] = 2.0; buf[2] = 3.0;
2754        buf[8] = 4.0; buf[9] = 5.0; buf[10] = 6.0;
2755        let a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
2756        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2757        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2758    }
2759
2760    // ****************************************************************
2761    // Matrix interop
2762    // ****************************************************************
2763
2764    #[cfg(feature = "matrix")]
2765    #[test]
2766    fn from_matrix() {
2767        let mat = Matrix::from_f64_unaligned(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, Some("m".into()));
2768        let a = NdArray::from(mat);
2769        assert_eq!(a.shape(), &[3, 2]);
2770        assert_eq!(a.name.as_deref(), Some("m"));
2771        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2772        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2773    }
2774
2775    #[cfg(feature = "matrix")]
2776    #[test]
2777    fn to_matrix_roundtrip() {
2778        let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2779        let a = NdArray::from_slice(&data, &[3, 2]);
2780        let mat = a.to_matrix().unwrap();
2781        assert_eq!(mat.n_rows, 3);
2782        assert_eq!(mat.n_cols, 2);
2783        assert_eq!(mat.col(0), &[1.0, 2.0, 3.0]);
2784        assert_eq!(mat.col(1), &[4.0, 5.0, 6.0]);
2785    }
2786
2787    #[cfg(feature = "matrix")]
2788    #[test]
2789    fn to_matrix_non_2d_fails() {
2790        let a = NdArray::new(&[5]);
2791        assert!(a.to_matrix().is_err());
2792    }
2793
2794    // ****************************************************************
2795    // Table interop
2796    // ****************************************************************
2797
2798    fn make_numeric_table() -> Table {
2799        let c0 = FieldArray::from_arr("x", Array::NumericArray(
2800            NumericArray::Float64(Arc::new(FloatArray::from_slice(&[1.0, 2.0, 3.0])))
2801        ));
2802        let c1 = FieldArray::from_arr("y", Array::NumericArray(
2803            NumericArray::Float64(Arc::new(FloatArray::from_slice(&[4.0, 5.0, 6.0])))
2804        ));
2805        Table::new("data".to_string(), Some(vec![c0, c1]))
2806    }
2807
2808    #[test]
2809    fn try_from_table() {
2810        let table = make_numeric_table();
2811        let a = NdArray::try_from(&table).unwrap();
2812        assert_eq!(a.shape(), &[3, 2]);
2813        assert_eq!(a.col(0), &[1.0, 2.0, 3.0]);
2814        assert_eq!(a.col(1), &[4.0, 5.0, 6.0]);
2815        assert_eq!(a.name.as_deref(), Some("data"));
2816    }
2817
2818    #[test]
2819    fn try_from_table_with_nulls_converts_to_nan() {
2820        let mut mask = Bitmask::new_set_all(3, true);
2821        mask.set(1, false);
2822        let arr = FloatArray::new(Buffer::from_slice(&[10.0, 0.0, 30.0]), Some(mask));
2823        let c0 = FieldArray::from_arr("v", Array::NumericArray(
2824            NumericArray::Float64(Arc::new(arr))
2825        ));
2826        let table = Table::new("nulls".to_string(), Some(vec![c0]));
2827        let a = NdArray::try_from(&table).unwrap();
2828        assert_eq!(a.get(&[0, 0]), 10.0);
2829        assert!(a.get(&[1, 0]).is_nan());
2830        assert_eq!(a.get(&[2, 0]), 30.0);
2831    }
2832
2833    #[test]
2834    fn try_from_table_coerces_unparseable_text_to_nan() {
2835        // TryFrom<&Table> uses the library's lenient numeric cast. Text that
2836        // does not parse as a number coerces to nulls, which surface as NaN
2837        // in the dense NdArray rather than failing the conversion.
2838        let c0 = FieldArray::from_arr("name", Array::from_string32(
2839            StringArray::from_slice(&["a", "b"])
2840        ));
2841        let table = Table::new("text".to_string(), Some(vec![c0]));
2842        let a = NdArray::try_from(&table).unwrap();
2843        assert_eq!(a.shape(), &[2, 1]);
2844        assert!(a.get(&[0, 0]).is_nan());
2845        assert!(a.get(&[1, 0]).is_nan());
2846    }
2847
2848    #[test]
2849    fn to_table_roundtrip() {
2850        let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2851        let a = NdArray::from_slice(&data, &[3, 2]);
2852        let fields = vec![
2853            Field::new("x", ArrowType::Float64, false, None),
2854            Field::new("y", ArrowType::Float64, false, None),
2855        ];
2856        let table = a.to_table(Some(fields)).unwrap();
2857        assert_eq!(table.n_rows(), 3);
2858        assert_eq!(table.n_cols(), 2);
2859        assert_eq!(table.col_names(), vec!["x", "y"]);
2860        let col0 = table.cols[0].array.num().f64();
2861        assert_eq!(col0.data.as_slice(), &[1.0, 2.0, 3.0]);
2862    }
2863
2864    #[test]
2865    fn to_table_generated_names() {
2866        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2867        let table = a.to_table(None).unwrap();
2868        assert_eq!(table.col_names(), vec!["col_0", "col_1"]);
2869    }
2870
2871    #[test]
2872    fn to_table_rejects_incompatible_field_metadata() {
2873        let wrong_dtype = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
2874            Field::new("value", ArrowType::Float32, false, None),
2875        ]));
2876        assert!(matches!(wrong_dtype, Err(MinarrowError::TypeError { .. })));
2877
2878        let nullable = NdArray::from_slice(&[1.0, 2.0], &[2, 1]).to_table(Some(vec![
2879            Field::new("value", ArrowType::Float64, true, None),
2880        ]));
2881        assert!(matches!(nullable, Err(MinarrowError::NullError { .. })));
2882    }
2883
2884    #[test]
2885    fn to_table_non_2d_fails() {
2886        let a = NdArray::new(&[5]);
2887        assert!(a.to_table(None).is_err());
2888    }
2889
2890    #[test]
2891    fn to_array_1d() {
2892        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2893        let arr = a.to_array().unwrap();
2894        let f = arr.num().f64();
2895        assert_eq!(f.data.as_slice(), &[1.0, 2.0, 3.0]);
2896    }
2897
2898    #[test]
2899    fn to_array_materialises_strided_logical_values() {
2900        let a = NdArray::from_buffer(
2901            Buffer::from_slice(&[1.0, 99.0, 2.0, 99.0, 3.0]),
2902            &[3],
2903            &[2],
2904        );
2905        let arr = a.to_array().unwrap();
2906        assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
2907    }
2908
2909    #[test]
2910    fn to_array_ignores_unused_backing_elements() {
2911        let a = NdArray::from_buffer(
2912            Buffer::from_slice(&[1.0, 2.0, 3.0, 99.0]),
2913            &[3],
2914            &[1],
2915        );
2916        let arr = a.to_array().unwrap();
2917        assert_eq!(arr.num().f64().data.as_slice(), &[1.0, 2.0, 3.0]);
2918    }
2919
2920    #[test]
2921    fn to_array_non_1d_fails() {
2922        let a = NdArray::new(&[2, 3]);
2923        assert!(a.to_array().is_err());
2924    }
2925
2926    // ****************************************************************
2927    // Concatenate
2928    // ****************************************************************
2929
2930    #[test]
2931    fn concat_1d() {
2932        let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2933        let b = NdArray::from_slice(&[3.0, 4.0, 5.0], &[3]);
2934        let c = a.concat(b).unwrap();
2935        assert_eq!(c.shape(), &[5]);
2936        let vals: Vec<f64> = (&c).into_iter().collect();
2937        assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
2938    }
2939
2940    #[test]
2941    fn concat_2d() {
2942        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2943        let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &[3, 2]);
2944        let c = a.concat(b).unwrap();
2945        assert_eq!(c.shape(), &[5, 2]);
2946        assert_eq!(c.col(0), &[1.0, 2.0, 5.0, 6.0, 7.0]);
2947        assert_eq!(c.col(1), &[3.0, 4.0, 8.0, 9.0, 10.0]);
2948    }
2949
2950    #[test]
2951    fn concat_dimension_mismatch_fails() {
2952        let a = NdArray::<f64>::new(&[3, 2]);
2953        let b = NdArray::new(&[3, 3]);
2954        assert!(a.concat(b).is_err());
2955    }
2956
2957    #[test]
2958    fn concat_rank_mismatch_fails() {
2959        let a = NdArray::<f64>::new(&[3]);
2960        let b = NdArray::new(&[3, 2]);
2961        assert!(a.concat(b).is_err());
2962    }
2963
2964    // ****************************************************************
2965    // Clone and PartialEq
2966    // ****************************************************************
2967
2968    #[test]
2969    fn clone_and_eq() {
2970        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2971        let b = a.clone();
2972        assert_eq!(a, b);
2973    }
2974
2975    #[test]
2976    fn ne_different_data() {
2977        let a = NdArray::from_slice(&[1.0, 2.0], &[2]);
2978        let b = NdArray::from_slice(&[1.0, 3.0], &[2]);
2979        assert_ne!(a, b);
2980    }
2981
2982    #[test]
2983    fn ne_different_shape() {
2984        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[4]);
2985        let b = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
2986        assert_ne!(a, b);
2987    }
2988
2989    // ****************************************************************
2990    // Debug formatting
2991    // ****************************************************************
2992
2993    #[test]
2994    fn debug_1d() {
2995        let a = NdArray::from_slice(&[1.0, 2.0, 3.0], &[3]);
2996        let s = format!("{:?}", a);
2997        assert!(s.contains("[3]"));
2998        assert!(s.contains("1D"));
2999    }
3000
3001    #[test]
3002    fn debug_2d_named() {
3003        let a = NdArray::<f64>::new_named(&[2, 3], "test");
3004        let s = format!("{:?}", a);
3005        assert!(s.contains("'test'"));
3006        assert!(s.contains("[2, 3]"));
3007        assert!(s.contains("2D"));
3008    }
3009
3010    // ****************************************************************
3011    // Edge cases
3012    // ****************************************************************
3013
3014    #[test]
3015    fn empty_array() {
3016        let a = NdArray::new(&[0, 5]);
3017        assert!(a.is_empty());
3018        assert_eq!(a.len(), 0);
3019        let vals: Vec<f64> = (&a).into_iter().collect();
3020        assert!(vals.is_empty());
3021    }
3022
3023    #[test]
3024    fn single_element() {
3025        let a = NdArray::from_slice(&[42.0], &[1]);
3026        assert_eq!(a.len(), 1);
3027        assert_eq!(a[(0,)], 42.0);
3028        let vals: Vec<f64> = (&a).into_iter().collect();
3029        assert_eq!(vals, vec![42.0]);
3030    }
3031
3032    #[test]
3033    fn single_element_2d() {
3034        let a = NdArray::from_slice(&[42.0], &[1, 1]);
3035        assert_eq!(a[(0, 0)], 42.0);
3036    }
3037
3038    #[test]
3039    fn large_array_iteration_count() {
3040        let n = 1000;
3041        let a = NdArray::<f64>::ones(&[n, 100]);
3042        let count = (&a).into_iter().count();
3043        assert_eq!(count, n * 100);
3044    }
3045
3046    // ****************************************************************
3047    // Bracket indexing: arr[col][row]
3048    // ****************************************************************
3049
3050    #[test]
3051    fn bracket_index_1d() {
3052        let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3053        assert_eq!(a[0], [10.0]);
3054        assert_eq!(a[2], [30.0]);
3055    }
3056
3057    #[test]
3058    fn bracket_index_2d_column() {
3059        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3060        assert_eq!(a[0], [1.0, 2.0, 3.0]);
3061        assert_eq!(a[1], [4.0, 5.0, 6.0]);
3062    }
3063
3064    #[test]
3065    fn bracket_index_2d_chained() {
3066        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3067        // arr[col][row]
3068        assert_eq!(a[0][0], 1.0);
3069        assert_eq!(a[0][2], 3.0);
3070        assert_eq!(a[1][0], 4.0);
3071        assert_eq!(a[1][2], 6.0);
3072    }
3073
3074    #[test]
3075    fn bracket_index_mut_2d() {
3076        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3077        a[0][1] = 99.0;
3078        assert_eq!(a[0][1], 99.0);
3079        assert_eq!(a[0], [1.0, 99.0, 3.0]);
3080    }
3081
3082    // ****************************************************************
3083    // Range indexing: arr[1..3]
3084    // ****************************************************************
3085
3086    #[test]
3087    fn range_index_1d() {
3088        let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3089        assert_eq!(a[1..3], [20.0, 30.0]);
3090    }
3091
3092    #[test]
3093    fn range_index_2d_columns() {
3094        // Selecting a range of columns returns the exact logical data,
3095        // since the compact layout has no padding between columns.
3096        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3097        let slab = &a[0..2];
3098        assert_eq!(slab, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3099        let col1 = &a[1..2];
3100        assert_eq!(col1, &[4.0, 5.0, 6.0]);
3101    }
3102
3103    #[cfg(feature = "matrix")]
3104    #[test]
3105    #[should_panic(expected = "contiguous")]
3106    fn range_index_non_contiguous_panics() {
3107        // A Matrix-imported array carries the padded stride, so range
3108        // indexing has no gap-free slab to return and panics.
3109        let mat = Matrix::from_f64_unaligned(
3110            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
3111        );
3112        let a = NdArray::from(mat);
3113        if a.is_contiguous() {
3114            // Padding only appears when rows are off the alignment boundary,
3115            // which holds for 3 rows. Guard the premise.
3116            panic!("premise failed: expected non-contiguous import");
3117        }
3118        let _ = &a[0..2];
3119    }
3120
3121    #[cfg(feature = "matrix")]
3122    #[test]
3123    fn to_matrix_repacks_compact_layout() {
3124        // Compact tensor data re-lays into Matrix's padded column layout.
3125        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3126        let mat = a.to_matrix().unwrap();
3127        assert_eq!(mat.n_rows, 3);
3128        assert_eq!(mat.n_cols, 2);
3129        assert_eq!(mat.stride, 8);
3130        assert_eq!(&mat.data.as_slice()[..3], &[1.0, 2.0, 3.0]);
3131        assert_eq!(&mat.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
3132    }
3133
3134    #[cfg(feature = "matrix")]
3135    #[test]
3136    fn matrix_roundtrip_via_contiguous() {
3137        // Matrix -> NdArray is zero-copy with the padded stride carried
3138        // through. to_contiguous compacts, and to_matrix re-pads.
3139        let mat = Matrix::from_f64_unaligned(
3140            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2, None,
3141        );
3142        let a = NdArray::from(mat);
3143        assert!(!a.is_contiguous());
3144        assert_eq!(a.get(&[2, 1]), 6.0);
3145        let compact = a.to_contiguous();
3146        assert!(compact.is_contiguous());
3147        assert_eq!(&compact[0..2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3148        let back = compact.to_matrix().unwrap();
3149        assert_eq!(back.stride, 8);
3150        assert_eq!(&back.data.as_slice()[8..11], &[4.0, 5.0, 6.0]);
3151    }
3152
3153    #[test]
3154    fn range_from_index() {
3155        let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3156        assert_eq!(a[2..], [30.0, 40.0]);
3157    }
3158
3159    #[test]
3160    fn range_to_index() {
3161        let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3162        assert_eq!(a[..2], [10.0, 20.0]);
3163    }
3164
3165    #[test]
3166    fn range_full_index() {
3167        let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3168        assert_eq!(a[..], [10.0, 20.0, 30.0]);
3169    }
3170
3171    // ****************************************************************
3172    // Slicing: arr.slice(nd![1..4, 2..5])
3173    // ****************************************************************
3174
3175    #[cfg(feature = "views")]
3176    #[test]
3177    fn slice_1d_single_index() {
3178        let a = NdArray::from_slice(&[10.0, 20.0, 30.0], &[3]);
3179        let v = a.slice(&[&1]);
3180        assert!(v.shape().is_empty());
3181        assert_eq!(v[()], 20.0);
3182    }
3183
3184    #[cfg(feature = "views")]
3185    #[test]
3186    fn slice_1d_range() {
3187        let a = NdArray::from_slice(&[10.0, 20.0, 30.0, 40.0], &[4]);
3188        let v = a.slice(&[&(1..3)]);
3189        assert_eq!(v.shape(), &[2]);
3190        assert_eq!(v[(0,)], 20.0);
3191        assert_eq!(v[(1,)], 30.0);
3192    }
3193
3194    #[cfg(feature = "views")]
3195    #[test]
3196    fn slice_2d_row_range_single_col() {
3197        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3198        // Rows 0..2 of column 1
3199        let v = a.slice(nd![0..2, 1]);
3200        assert_eq!(v.shape(), &[2]);
3201        assert_eq!(v[(0,)], 4.0);
3202        assert_eq!(v[(1,)], 5.0);
3203    }
3204
3205    #[cfg(feature = "views")]
3206    #[test]
3207    fn slice_2d_single_row_col_range() {
3208        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3209        // Row 1, columns 0..2 - collapses row axis
3210        let v = a.slice(nd![1, 0..2]);
3211        assert_eq!(v.shape(), &[2]);
3212        // Should get row 1 values: a[(1,0)]=2.0, a[(1,1)]=5.0
3213        let vals: Vec<f64> = (&v).into_iter().collect();
3214        assert_eq!(vals, vec![2.0, 5.0]);
3215    }
3216
3217    #[cfg(feature = "views")]
3218    #[test]
3219    fn slice_2d_both_ranges() {
3220        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3221        // Rows 0..2, columns 0..2 - sub-matrix
3222        let v = a.slice(nd![0..2, 0..2]);
3223        assert_eq!(v.shape(), &[2, 2]);
3224        assert_eq!(v[(0, 0)], 1.0);
3225        assert_eq!(v[(1, 0)], 2.0);
3226        assert_eq!(v[(0, 1)], 4.0);
3227        assert_eq!(v[(1, 1)], 5.0);
3228    }
3229
3230    #[cfg(feature = "views")]
3231    #[test]
3232    fn slice_2d_both_indices_scalar() {
3233        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3234        // Single element as a rank-zero scalar view
3235        let v = a.slice(nd![2, 1]);
3236        assert!(v.shape().is_empty());
3237        assert_eq!(v[()], 6.0);
3238    }
3239
3240    #[cfg(feature = "views")]
3241    #[test]
3242    fn slice_3d_mixed() {
3243        // 2x3x4 array
3244        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3245        let a = NdArray::from_slice(&data, &[2, 3, 4]);
3246        // All rows, column 1, slices 0..2
3247        let v = a.slice(nd![0..2, 1, 0..2]);
3248        assert_eq!(v.shape(), &[2, 2]);
3249        // a[(0,1,0)]=3, a[(1,1,0)]=4, a[(0,1,1)]=9, a[(1,1,1)]=10
3250        assert_eq!(v[(0, 0)], 3.0);
3251        assert_eq!(v[(1, 0)], 4.0);
3252        assert_eq!(v[(0, 1)], 9.0);
3253        assert_eq!(v[(1, 1)], 10.0);
3254    }
3255
3256    #[cfg(feature = "views")]
3257    #[test]
3258    fn slice_with_nd_macro() {
3259        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3260        let a = NdArray::from_slice(&data, &[2, 3, 4]);
3261        let v = a.slice(nd![0..2, 0..3, 0..4]);
3262        assert_eq!(v.shape(), &[2, 3, 4]);
3263        assert_eq!(v.len(), 24);
3264    }
3265
3266    #[cfg(feature = "views")]
3267    #[test]
3268    fn slice_preserves_data_through_iteration() {
3269        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3270        let v = a.slice(nd![1..3, 0..2]);
3271        // Sub-matrix: rows 1..3, cols 0..2
3272        let vals: Vec<f64> = (&v).into_iter().collect();
3273        assert_eq!(vals, vec![2.0, 3.0, 5.0, 6.0]);
3274    }
3275
3276    #[cfg(feature = "views")]
3277    #[test]
3278    fn slice_column_window() {
3279        // Slice rows 2..5 from column 1 of a [10, 2] array
3280        let data: Vec<f64> = (1..=20).map(|x| x as f64).collect();
3281        let a = NdArray::from_slice(&data, &[10, 2]);
3282        let v = a.slice(nd![2..5, 1]);
3283        assert_eq!(v.shape(), &[3]);
3284        assert_eq!(v[(0,)], 13.0);
3285        assert_eq!(v[(1,)], 14.0);
3286        assert_eq!(v[(2,)], 15.0);
3287    }
3288
3289    // ****************************************************************
3290    // Bounds and panic contracts
3291    // ****************************************************************
3292
3293    #[test]
3294    #[should_panic(expected = "Column index out of bounds")]
3295    fn col_out_of_bounds_panics() {
3296        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3297        let _ = a.col(2);
3298    }
3299
3300    #[test]
3301    #[should_panic(expected = "indices for a 2D array")]
3302    fn get_rank_mismatch_panics() {
3303        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3304        let _ = a.get(&[0]);
3305    }
3306
3307    #[test]
3308    #[should_panic(expected = "out of bounds for dim")]
3309    fn get_index_out_of_bounds_panics() {
3310        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3311        let _ = a.get(&[2, 0]);
3312    }
3313
3314    #[test]
3315    #[should_panic(expected = "out of bounds for dim")]
3316    fn set_out_of_bounds_panics() {
3317        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3318        a.set(&[0, 5], 9.0);
3319    }
3320
3321    #[cfg(all(feature = "views", feature = "select"))]
3322    #[test]
3323    #[should_panic(expected = "expected 2 axes, got 1")]
3324    fn slice_wrong_axis_count_panics() {
3325        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3326        let _ = a.slice(nd![0..2]);
3327    }
3328
3329    #[cfg(all(feature = "views", feature = "select"))]
3330    #[test]
3331    #[should_panic(expected = "range 0..100 out of bounds")]
3332    fn slice_range_out_of_bounds_panics() {
3333        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3334        let _ = a.slice(nd![0..100, 0..2]);
3335    }
3336
3337    #[cfg(all(feature = "views", feature = "select"))]
3338    #[test]
3339    #[should_panic(expected = "index 5 out of bounds")]
3340    fn slice_single_index_out_of_bounds_panics() {
3341        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3342        let _ = a.slice(nd![5, 0..2]);
3343    }
3344
3345    #[test]
3346    #[should_panic(expected = "at least 2 points")]
3347    fn linspace_requires_two_points() {
3348        let _ = NdArray::<f64>::linspace(0.0, 1.0, 1);
3349    }
3350
3351    #[cfg(all(feature = "views", feature = "select"))]
3352    #[test]
3353    fn slice_one_element_index_array_collapses() {
3354        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3355        let idx: &[usize] = &[1];
3356        let v = a.slice(nd![idx, 0..2]);
3357        assert_eq!(v.shape(), &[2]);
3358        assert_eq!(v[(0,)], 2.0);
3359        assert_eq!(v[(1,)], 5.0);
3360    }
3361
3362    #[test]
3363    fn get_unchecked_matches_get() {
3364        let data: Vec<f64> = (1..=24).map(|x| x as f64).collect();
3365        let a = NdArray::from_slice(&data, &[2, 3, 4]);
3366        for i in 0..2 {
3367            for j in 0..3 {
3368                for k in 0..4 {
3369                    let idx = [i, j, k];
3370                    // SAFETY: indices are within shape
3371                    assert_eq!(unsafe { a.get_unchecked(&idx) }, a.get(&idx));
3372                }
3373            }
3374        }
3375    }
3376
3377    // ****************************************************************
3378    // Concatenate correctness
3379    // ****************************************************************
3380
3381    #[test]
3382    fn concat_3d_interleaves_axis0() {
3383        let da: Vec<f64> = (1..=8).map(|x| x as f64).collect();
3384        let db: Vec<f64> = (9..=16).map(|x| x as f64).collect();
3385        let a = NdArray::from_slice(&da, &[2, 2, 2]);
3386        let b = NdArray::from_slice(&db, &[2, 2, 2]);
3387        let c = a.clone().concat(b.clone()).unwrap();
3388        assert_eq!(c.shape(), &[4, 2, 2]);
3389        for i in 0..4 {
3390            for j in 0..2 {
3391                for k in 0..2 {
3392                    let expected = if i < 2 {
3393                        a.get(&[i, j, k])
3394                    } else {
3395                        b.get(&[i - 2, j, k])
3396                    };
3397                    assert_eq!(c.get(&[i, j, k]), expected);
3398                }
3399            }
3400        }
3401    }
3402
3403    #[test]
3404    fn concat_non_contiguous_operand() {
3405        // Row-major strides on the first operand exercise the general
3406        // interleave path.
3407        let a = NdArray::from_buffer(
3408            Buffer::from_slice(&[1.0, 2.0, 3.0, 4.0]),
3409            &[2, 2],
3410            &[2, 1],
3411        );
3412        assert!(!a.is_contiguous());
3413        let b = NdArray::from_slice(&[5.0, 6.0, 7.0, 8.0], &[2, 2]);
3414        let c = a.concat(b).unwrap();
3415        assert_eq!(c.shape(), &[4, 2]);
3416        assert_eq!(c.col(0), &[1.0, 3.0, 5.0, 6.0]);
3417        assert_eq!(c.col(1), &[2.0, 4.0, 7.0, 8.0]);
3418    }
3419
3420    // ****************************************************************
3421    // Copy-on-write and strided mutation
3422    // ****************************************************************
3423
3424    #[cfg(feature = "views")]
3425    #[test]
3426    fn set_after_view_copy_on_write() {
3427        let mut a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
3428        let v = a.as_view();
3429        a.set(&[0, 0], 99.0);
3430        // The write detached the array's buffer, so the view keeps the
3431        // original value.
3432        assert_eq!(a.get(&[0, 0]), 99.0);
3433        assert_eq!(v.get(&[0, 0]), 1.0);
3434    }
3435
3436    #[test]
3437    fn fill_with_non_contiguous_logical_only() {
3438        // Padded column stride leaves gaps between columns in the buffer.
3439        let mut buf = Vec64::with_capacity(11);
3440        buf.0.resize(11, 0.0);
3441        let mut a = NdArray::from_buffer(Buffer::from_vec64(buf), &[3, 2], &[1, 8]);
3442        assert!(!a.is_contiguous());
3443        a.fill_with(7.0);
3444        for v in &a { assert_eq!(v, 7.0); }
3445        // Padding between the columns stays untouched.
3446        assert_eq!(a.as_slice()[3], 0.0);
3447        assert_eq!(a.as_slice()[7], 0.0);
3448    }
3449
3450    // ****************************************************************
3451    // Range shorthand on 2D
3452    // ****************************************************************
3453
3454    #[test]
3455    fn range_from_index_2d() {
3456        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3457        assert_eq!(&a[1..], &[4.0, 5.0, 6.0]);
3458    }
3459
3460    #[test]
3461    fn range_to_index_2d() {
3462        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3463        assert_eq!(&a[..1], &[1.0, 2.0, 3.0]);
3464    }
3465
3466    #[test]
3467    fn range_full_index_2d() {
3468        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3469        assert_eq!(&a[..], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
3470    }
3471
3472    // ****************************************************************
3473    // Row window and size estimate
3474    // ****************************************************************
3475
3476    #[cfg(all(feature = "views", feature = "select"))]
3477    #[test]
3478    fn row_selection_single_index_window() {
3479        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3480        assert_eq!(a.n_obs(), 3);
3481        let v = a.r(1usize);
3482        assert_eq!(v.shape(), &[1, 2]);
3483        assert_eq!(v.get(&[0, 0]), 2.0);
3484        assert_eq!(v.get(&[0, 1]), 5.0);
3485        // A window keeps the full source buffer rather than gathering a copy.
3486        assert_eq!(v.source.len(), a.len());
3487        // SAFETY: indices are within the window shape
3488        assert_eq!(unsafe { v.get_unchecked(&[0, 1]) }, 5.0);
3489    }
3490
3491    #[cfg(feature = "size")]
3492    #[test]
3493    fn est_bytes_covers_buffer() {
3494        use crate::traits::byte_size::ByteSize;
3495        let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
3496        assert!(a.est_bytes() >= 6 * size_of::<f64>());
3497    }
3498}