minarrow/traits/
selection.rs

1// Copyright 2025 Peter Garfield Bower
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! # **Selection Traits** - *Selection across dimensions*
16//!
17//! Traits for field and data selection that enable polymorphic methods
18//! across Table, TableV, Cube, and future types.
19//!
20//! ## Architecture
21//! - **FieldSelector**: Input types that can specify field selection (e.g., `&str`, `usize`)
22//! - **DataSelector**: Input types that can specify data selection (e.g., `usize`, ranges)
23//! - **FieldSelection**: Capability trait for types that support field selection
24//! - **DataSelection**: Capability trait for types that support data selection
25//! - **Selection2D**: Combined 2D selection (FieldSelection + DataSelection)
26//! - **Selection3D**: Future Extension for 3D selection
27//! - **Selection4D**: Future Extension for 4D selection
28
29use crate::Field;
30use std::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo};
31use std::sync::Arc;
32
33// Input types that can be passed to selection methods
34// ===================================================
35// These traits are implemented on user-facing input types like `&str`, `usize`, and ranges.
36// They convert user input (e.g., `table.f("name")` or `table.d(0..10)`) into index vectors.
37// These are "what the user writes" when selecting.
38
39/// Trait for types that can specify a field selection (named/schema dimension)
40pub trait FieldSelector {
41    /// Resolve this selection to field indices for the given fields
42    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize>;
43
44    /// Produce an owned version of this selector.
45    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync>;
46}
47
48/// Trait for types that can specify a data selection (index-based dimension)
49pub trait DataSelector {
50    /// Resolve this selection to indices (within the given count)
51    fn resolve_indices(&self, count: usize) -> Vec<usize>;
52
53    /// Returns true if this selector represents a contiguous range.
54    /// Range types (Range, RangeFrom, etc.) return true.
55    /// Index arrays (&[usize], Vec<usize>) return false.
56    fn is_contiguous(&self) -> bool {
57        false // Default: assume non-contiguous
58    }
59
60    /// Resolve this selector against one axis of `dim_size`, producing a
61    /// `(start, end, collapse)` span. Ranges window the axis and a single
62    /// index collapses it. Panics when the selection falls outside the
63    /// axis. Index arrays have no strided-view representation, so they
64    /// are rejected - gather through [`RowSelection::r`] instead.
65    #[cfg(feature = "ndarray")]
66    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
67        // usize::MAX so out-of-range indices survive to the bounds check
68        // below rather than being silently filtered.
69        let indices = self.resolve_indices(usize::MAX);
70        assert_eq!(
71            indices.len(), 1,
72            "axis selection: index arrays take a single index or a contiguous range; gather with r() instead"
73        );
74        assert!(
75            indices[0] < dim_size,
76            "axis selection: index {} out of bounds (size {})", indices[0], dim_size
77        );
78        (indices[0], indices[0] + 1, true)
79    }
80}
81
82// These traits are implemented on structures like Table, ArrayV, etc.
83// They define what selection methods are available on each structure.
84// These are "what the structure can do" for selection operations.
85
86/// Trait for types that support field/column selection
87///
88/// Associated types determine what each access pattern returns:
89/// - `View`: multi-field selection result, e.g. TableV for Table, Cube for Cube
90/// - `DataView`: single column by index, e.g. ArrayV
91/// - `Field`: single field by name via `.get()`, e.g. FieldArray for Table, Arc<Table> for Cube
92pub trait ColumnSelection {
93    /// The view type returned by multi-field selection e.g. TableV
94    type View;
95    /// A single column view by index e.g. ArrayV
96    type ColumnView;
97    /// An owned single field by name via `.get()` e.g. FieldArray for Table, Arc<Table> for Cube
98    type ColumnOwned;
99
100    /// Select fields/columns by name, index, or range
101    ///
102    /// Shorthand alias for `col`
103    ///
104    /// # Examples
105    /// ```ignore
106    /// table.c("age")           // single column by name
107    /// table.c(&["a", "b"])     // multiple columns by name
108    /// table.c(0)               // single column by index
109    /// table.c(0..3)            // columns by range
110    /// ```
111    fn c<S: FieldSelector>(&self, selection: S) -> Self::View;
112
113    /// Select a single column by name
114    ///
115    /// Named form of `c` - use `c` for selection by index, range,
116    /// or multiple names.
117    ///
118    /// # Examples
119    /// ```ignore
120    /// table.col("age")         // single column by name
121    /// table.c(&["a", "b"])     // multiple columns by name
122    /// table.c(0)               // single column by index
123    /// table.c(0..3)            // columns by range
124    /// ```
125    fn col(&self, name: &str) -> Self::View {
126        self.c(name)
127    }
128
129    /// Get a single field by name, returning an owned value.
130    fn get(&self, field: &str) -> Option<Self::ColumnOwned>;
131
132    /// Get a single column view by index
133    fn col_ix(&self, idx: usize) -> Option<Self::ColumnView>;
134
135    /// Get all columns as views
136    fn col_vec(&self) -> Vec<Self::ColumnView>;
137
138    /// Get the fields for field resolution
139    fn get_cols(&self) -> Vec<Arc<Field>>;
140}
141
142/// Trait for types that support row/data selection
143pub trait RowSelection {
144    /// The view type returned by selection operations
145    type View;
146
147    /// Select rows by index or range
148    ///
149    /// Shorthand alias for `row`
150    ///
151    /// # Examples
152    /// ```ignore
153    /// table.r(5)               // single row
154    /// table.r(&[1, 3, 5])      // specific rows
155    /// table.r(0..10)           // row range
156    /// ```
157    fn r<S: DataSelector>(&self, selection: S) -> Self::View;
158
159    /// Select a single row by index
160    ///
161    /// Named form of `r` - use `r` for selection by range or
162    /// index array.
163    ///
164    /// # Examples
165    /// ```ignore
166    /// table.row(5)             // single row
167    /// table.r(&[1, 3, 5])      // specific rows
168    /// table.r(0..10)           // row range
169    /// ```
170    fn row(&self, idx: usize) -> Self::View {
171        self.r(idx)
172    }
173
174    /// Get the count for data resolution
175    fn get_row_count(&self) -> usize;
176}
177
178/// Trait for types that support selection across any axis combination.
179///
180/// The N-dimensional member of the selection family, alongside
181/// [`ColumnSelection`] and [`RowSelection`]. Each axis takes any
182/// [`DataSelector`] - a single index collapses the dimension, and a
183/// contiguous range keeps it.
184///
185/// # Examples
186/// ```ignore
187/// arr.s(&[&(1..4), &2])        // rows 1..4 of column 2 (2D)
188/// arr.s(nd![1..4, 2])          // the same through the nd! macro
189/// arr.s(nd![0..2, 1, 0..3])    // mixed selection (3D)
190/// ```
191#[cfg(feature = "ndarray")]
192pub trait AxisSelection {
193    /// The view type returned by axis selection e.g. NdArrayV
194    type View;
195
196    /// Select along every axis at once, one [`DataSelector`] per axis
197    ///
198    /// Shorthand alias for `select`
199    ///
200    /// A single index collapses its dimension, and a contiguous range
201    /// keeps it.
202    ///
203    /// # Examples
204    /// ```ignore
205    /// arr.s(nd![1..4, 2])          // rows 1..4 of column 2 (2D)
206    /// arr.s(nd![0..2, 1, 0..3])    // mixed selection (3D)
207    /// arr.s(nd![.., 5])            // full range on axis 0
208    /// arr.s(&[&(1..4), &2])        // without the nd! macro
209    /// ```
210    fn s(&self, selection: &[&dyn DataSelector]) -> Self::View;
211
212    /// Select along every axis at once, one [`DataSelector`] per axis
213    ///
214    /// A single index collapses its dimension, and a contiguous range
215    /// keeps it.
216    ///
217    /// # Examples
218    /// ```ignore
219    /// arr.select(nd![1..4, 2])          // rows 1..4 of column 2 (2D)
220    /// arr.select(nd![0..2, 1, 0..3])    // mixed selection (3D)
221    /// arr.select(nd![.., 5])            // full range on axis 0
222    /// ```
223    fn select(&self, selection: &[&dyn DataSelector]) -> Self::View {
224        self.s(selection)
225    }
226
227    /// Get the axis count for selection resolution
228    fn get_axis_count(&self) -> usize;
229}
230
231/// Combined trait for 2D selection (field + data dimensions)
232///
233/// This trait is automatically implemented for any type that implements
234/// both `ColumnSelection` and `RowSelection` with the same `View` type.
235pub trait Selection2D: ColumnSelection + RowSelection {}
236
237/// Blanket implementation for any type that implements both traits
238impl<T> Selection2D for T where
239    T: ColumnSelection + RowSelection<View = <T as ColumnSelection>::View>
240{
241}
242
243// These allow users to pass names, indices, and ranges when selecting fields.
244// For example: table.c("age"), table.c(&["name", "age"]), table.c(0..3)
245
246/// Single field by name
247impl FieldSelector for &str {
248    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
249        fields
250            .iter()
251            .position(|f| f.name == *self)
252            .into_iter()
253            .collect()
254    }
255    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
256        Box::new(vec![self.to_string()])
257    }
258}
259
260/// Multiple fields by names
261impl FieldSelector for &[&str] {
262    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
263        self.iter()
264            .filter_map(|name| fields.iter().position(|f| f.name == *name))
265            .collect()
266    }
267    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
268        Box::new(self.iter().map(|s| s.to_string()).collect::<Vec<String>>())
269    }
270}
271
272/// Multiple fields by names (array reference)
273impl<const N: usize> FieldSelector for &[&str; N] {
274    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
275        self.iter()
276            .filter_map(|name| fields.iter().position(|f| f.name == *name))
277            .collect()
278    }
279    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
280        Box::new(self.iter().map(|s| s.to_string()).collect::<Vec<String>>())
281    }
282}
283
284/// Multiple fields by names (Vec of borrowed str)
285impl FieldSelector for Vec<&str> {
286    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
287        self.iter()
288            .filter_map(|name| fields.iter().position(|f| f.name == *name))
289            .collect()
290    }
291    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
292        Box::new(self.iter().map(|s| s.to_string()).collect::<Vec<String>>())
293    }
294}
295
296/// Multiple fields by names (Vec of owned String)
297impl FieldSelector for Vec<String> {
298    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
299        self.iter()
300            .filter_map(|name| fields.iter().position(|f| f.name.as_str() == name.as_str()))
301            .collect()
302    }
303    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
304        Box::new(self.clone())
305    }
306}
307
308/// Single field by index
309impl FieldSelector for usize {
310    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
311        if *self < fields.len() {
312            vec![*self]
313        } else {
314            Vec::new()
315        }
316    }
317    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
318        Box::new(vec![*self])
319    }
320}
321
322/// Multiple fields by indices
323impl FieldSelector for &[usize] {
324    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
325        self.iter()
326            .copied()
327            .filter(|&idx| idx < fields.len())
328            .collect()
329    }
330    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
331        Box::new(self.to_vec())
332    }
333}
334
335/// Multiple fields by indices (array reference)
336impl<const N: usize> FieldSelector for &[usize; N] {
337    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
338        self.iter()
339            .copied()
340            .filter(|&idx| idx < fields.len())
341            .collect()
342    }
343    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
344        Box::new(self.to_vec())
345    }
346}
347
348/// Multiple fields by indices (Vec)
349impl FieldSelector for Vec<usize> {
350    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
351        self.iter()
352            .copied()
353            .filter(|&idx| idx < fields.len())
354            .collect()
355    }
356    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
357        Box::new(self.clone())
358    }
359}
360
361/// Field range selection
362impl FieldSelector for Range<usize> {
363    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
364        let end = self.end.min(fields.len());
365        (self.start..end).collect()
366    }
367    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
368        Box::new(self.clone())
369    }
370}
371
372/// Field range from selection
373impl FieldSelector for RangeFrom<usize> {
374    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
375        (self.start..fields.len()).collect()
376    }
377    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
378        Box::new(self.clone())
379    }
380}
381
382/// Field range to selection
383impl FieldSelector for RangeTo<usize> {
384    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
385        let end = self.end.min(fields.len());
386        (0..end).collect()
387    }
388    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
389        Box::new(self.clone())
390    }
391}
392
393/// Field full range selection
394impl FieldSelector for RangeFull {
395    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
396        (0..fields.len()).collect()
397    }
398    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
399        Box::new(..)
400    }
401}
402
403/// Field inclusive range selection
404impl FieldSelector for RangeInclusive<usize> {
405    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
406        let start = *self.start();
407        let end = (*self.end() + 1).min(fields.len());
408        (start..end).collect()
409    }
410    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
411        Box::new(self.clone())
412    }
413}
414
415/// Boxed field selector for owned selection
416impl FieldSelector for Box<dyn FieldSelector + Send + Sync> {
417    fn resolve_fields(&self, fields: &[Arc<Field>]) -> Vec<usize> {
418        (**self).resolve_fields(fields)
419    }
420    fn to_owned(&self) -> Box<dyn FieldSelector + Send + Sync> {
421        (**self).to_owned()
422    }
423}
424
425// These allow users to pass indices and ranges when selecting data (rows, time, etc.).
426// For example: table.r(5), table.r(&[1, 3, 5]), table.r(0..10)
427
428/// Single data index
429impl DataSelector for usize {
430    fn resolve_indices(&self, count: usize) -> Vec<usize> {
431        if *self < count {
432            vec![*self]
433        } else {
434            Vec::new()
435        }
436    }
437
438    /// A single index is a contiguous window of length one.
439    fn is_contiguous(&self) -> bool {
440        true
441    }
442
443    #[cfg(feature = "ndarray")]
444    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
445        assert!(
446            *self < dim_size,
447            "axis selection: index {} out of bounds (size {})", self, dim_size
448        );
449        (*self, *self + 1, true)
450    }
451}
452
453/// Single data index from a plain integer literal. Negative values
454/// resolve to nothing.
455impl DataSelector for i32 {
456    fn resolve_indices(&self, count: usize) -> Vec<usize> {
457        if *self >= 0 && (*self as usize) < count {
458            vec![*self as usize]
459        } else {
460            Vec::new()
461        }
462    }
463
464    /// A single index is a contiguous window of length one.
465    fn is_contiguous(&self) -> bool {
466        true
467    }
468
469    #[cfg(feature = "ndarray")]
470    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
471        assert!(
472            *self >= 0 && (*self as usize) < dim_size,
473            "axis selection: index {} out of bounds (size {})", self, dim_size
474        );
475        (*self as usize, *self as usize + 1, true)
476    }
477}
478
479/// Multiple data indices
480impl DataSelector for &[usize] {
481    fn resolve_indices(&self, count: usize) -> Vec<usize> {
482        self.iter().copied().filter(|&idx| idx < count).collect()
483    }
484}
485
486/// Multiple data indices (array reference)
487impl<const N: usize> DataSelector for &[usize; N] {
488    fn resolve_indices(&self, count: usize) -> Vec<usize> {
489        self.iter().copied().filter(|&idx| idx < count).collect()
490    }
491}
492
493/// Multiple data indices (Vec)
494impl DataSelector for Vec<usize> {
495    fn resolve_indices(&self, count: usize) -> Vec<usize> {
496        self.iter().copied().filter(|&idx| idx < count).collect()
497    }
498}
499
500/// Data range selection
501impl DataSelector for Range<usize> {
502    fn resolve_indices(&self, count: usize) -> Vec<usize> {
503        let end = self.end.min(count);
504        (self.start..end).collect()
505    }
506
507    fn is_contiguous(&self) -> bool {
508        true
509    }
510
511    #[cfg(feature = "ndarray")]
512    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
513        assert!(
514            self.start <= self.end && self.end <= dim_size,
515            "axis selection: range {}..{} out of bounds (size {})",
516            self.start, self.end, dim_size
517        );
518        (self.start, self.end, false)
519    }
520}
521
522/// Data range selection from plain integer literals. Negative bounds
523/// clamp to zero.
524impl DataSelector for Range<i32> {
525    fn resolve_indices(&self, count: usize) -> Vec<usize> {
526        let start = self.start.max(0) as usize;
527        let end = (self.end.max(0) as usize).min(count);
528        (start..end).collect()
529    }
530
531    fn is_contiguous(&self) -> bool {
532        true
533    }
534
535    #[cfg(feature = "ndarray")]
536    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
537        assert!(
538            self.start >= 0 && self.end >= self.start && (self.end as usize) <= dim_size,
539            "axis selection: range {}..{} out of bounds (size {})",
540            self.start, self.end, dim_size
541        );
542        (self.start as usize, self.end as usize, false)
543    }
544}
545
546/// Data range from selection
547impl DataSelector for RangeFrom<usize> {
548    fn resolve_indices(&self, count: usize) -> Vec<usize> {
549        (self.start..count).collect()
550    }
551
552    fn is_contiguous(&self) -> bool {
553        true
554    }
555
556    #[cfg(feature = "ndarray")]
557    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
558        assert!(
559            self.start <= dim_size,
560            "axis selection: range {}.. out of bounds (size {})", self.start, dim_size
561        );
562        (self.start, dim_size, false)
563    }
564}
565
566/// Data range from selection from a plain integer literal.
567impl DataSelector for RangeFrom<i32> {
568    fn resolve_indices(&self, count: usize) -> Vec<usize> {
569        let start = self.start.max(0) as usize;
570        (start..count).collect()
571    }
572
573    fn is_contiguous(&self) -> bool {
574        true
575    }
576
577    #[cfg(feature = "ndarray")]
578    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
579        assert!(
580            self.start >= 0 && (self.start as usize) <= dim_size,
581            "axis selection: range {}.. out of bounds (size {})", self.start, dim_size
582        );
583        (self.start as usize, dim_size, false)
584    }
585}
586
587/// Data range to selection
588impl DataSelector for RangeTo<usize> {
589    fn resolve_indices(&self, count: usize) -> Vec<usize> {
590        let end = self.end.min(count);
591        (0..end).collect()
592    }
593
594    fn is_contiguous(&self) -> bool {
595        true
596    }
597
598    #[cfg(feature = "ndarray")]
599    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
600        assert!(
601            self.end <= dim_size,
602            "axis selection: range ..{} out of bounds (size {})", self.end, dim_size
603        );
604        (0, self.end, false)
605    }
606}
607
608/// Data range to selection from a plain integer literal. Negative bounds
609/// clamp to zero.
610impl DataSelector for RangeTo<i32> {
611    fn resolve_indices(&self, count: usize) -> Vec<usize> {
612        let end = (self.end.max(0) as usize).min(count);
613        (0..end).collect()
614    }
615
616    fn is_contiguous(&self) -> bool {
617        true
618    }
619
620    #[cfg(feature = "ndarray")]
621    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
622        assert!(
623            self.end >= 0 && (self.end as usize) <= dim_size,
624            "axis selection: range ..{} out of bounds (size {})", self.end, dim_size
625        );
626        (0, self.end as usize, false)
627    }
628}
629
630/// Data full range selection
631impl DataSelector for RangeFull {
632    fn resolve_indices(&self, count: usize) -> Vec<usize> {
633        (0..count).collect()
634    }
635
636    fn is_contiguous(&self) -> bool {
637        true
638    }
639
640    #[cfg(feature = "ndarray")]
641    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
642        (0, dim_size, false)
643    }
644}
645
646/// Data inclusive range selection
647impl DataSelector for RangeInclusive<usize> {
648    fn resolve_indices(&self, count: usize) -> Vec<usize> {
649        let start = *self.start();
650        let end = (*self.end() + 1).min(count);
651        (start..end).collect()
652    }
653
654    fn is_contiguous(&self) -> bool {
655        true
656    }
657
658    #[cfg(feature = "ndarray")]
659    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
660        assert!(
661            self.start() <= self.end() && *self.end() < dim_size,
662            "axis selection: range {}..={} out of bounds (size {})",
663            self.start(), self.end(), dim_size
664        );
665        (*self.start(), *self.end() + 1, false)
666    }
667}
668
669/// Data inclusive range selection from plain integer literals. Negative
670/// bounds clamp to zero.
671impl DataSelector for RangeInclusive<i32> {
672    fn resolve_indices(&self, count: usize) -> Vec<usize> {
673        if *self.end() < 0 {
674            return Vec::new();
675        }
676        let start = (*self.start()).max(0) as usize;
677        let end = (*self.end() as usize + 1).min(count);
678        (start..end).collect()
679    }
680
681    fn is_contiguous(&self) -> bool {
682        true
683    }
684
685    #[cfg(feature = "ndarray")]
686    fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
687        assert!(
688            *self.start() >= 0 && self.start() <= self.end() && (*self.end() as usize) < dim_size,
689            "axis selection: range {}..={} out of bounds (size {})",
690            self.start(), self.end(), dim_size
691        );
692        (*self.start() as usize, *self.end() as usize + 1, false)
693    }
694}
minarrow/traits/selection.rs

minarrow/traits/
selection.rs
