Skip to main content

volas_core/
index.rs

1//! Index: the row labels shared by a frame and the series drawn from it.
2
3use crate::column::Column;
4use crate::error::{Result, VolasError};
5use crate::tz::Tz;
6
7/// Row labels plus an optional name.
8///
9/// `name` is index-wide metadata (pandas `Index.name`): it is recorded by
10/// `set_index` (from the source column), restored by `reset_index`, and shown in
11/// a frame's / series' repr. It is `None` for an unnamed index (the default
12/// `Range` index, a freshly built datetime index, …).
13///
14/// `kind` is the label storage. A `Datetime` kind carries its own [`Tz`]:
15/// storage is always UTC epoch-ns, but the tz governs how those instants render
16/// and how bare-string / day-bucket matching maps to wall-clock time (see
17/// [`crate::tz`]). Both the name and a datetime kind's tz ride with the shared
18/// `Arc<Index>`, so a frame and every series drawn from it agree on them for free.
19#[derive(Clone, Debug, PartialEq)]
20pub struct Index {
21    /// Optional index name (pandas `Index.name`); `None` when unnamed.
22    pub name: Option<String>,
23    /// The label storage / kind.
24    pub kind: IndexKind,
25}
26
27/// The label storage backing an [`Index`].
28///
29/// Defaults to an implicit `0..n` range; a `Datetime` kind is the common OHLCV
30/// case (i64 nanoseconds since the Unix epoch); a `Str` kind (pandas
31/// object/string index) supports symbol-keyed lookup.
32#[derive(Clone, Debug, PartialEq)]
33pub enum IndexKind {
34    /// Implicit `0..n` integer labels.
35    Range(usize),
36    /// Explicit integer labels.
37    Int64(Vec<i64>),
38    /// Datetime labels as i64 nanoseconds since the Unix epoch, with a display /
39    /// matching timezone (UTC by default).
40    Datetime(Vec<i64>, Tz),
41    /// String labels (pandas object/string index).
42    Str(Vec<String>),
43}
44
45impl IndexKind {
46    /// Materialize the numeric labels as `i64`. Numeric kinds only — a string
47    /// kind has no i64 labels (its callers guard against it).
48    fn to_i64_labels(&self) -> Vec<i64> {
49        match self {
50            IndexKind::Range(n) => (0..*n as i64).collect(),
51            IndexKind::Int64(v) => v.clone(),
52            IndexKind::Datetime(v, _) => v.clone(),
53            IndexKind::Str(_) => unreachable!("string indexes have no i64 labels"), // LCOV_EXCL_LINE
54        }
55    }
56}
57
58/// A single row-index label: an integer / datetime (`I64`, ns) or a string
59/// (`Str`). It decouples label *lookup* from the index's storage so the same
60/// `.loc` / `.at` / `.loc[a:b]` / `drop` paths serve every index kind.
61#[derive(Clone, Debug, PartialEq)]
62pub enum Label {
63    /// An integer or datetime-ns label.
64    I64(i64),
65    /// A string label.
66    Str(String),
67}
68
69impl Label {
70    /// The i64 payload, if this is an `I64` label.
71    pub fn as_i64(&self) -> Option<i64> {
72        if let Label::I64(v) = self {
73            Some(*v)
74        } else {
75            None
76        }
77    }
78
79    /// The string payload, if this is a `Str` label.
80    pub fn as_str(&self) -> Option<&str> {
81        if let Label::Str(s) = self {
82            Some(s.as_str())
83        } else {
84            None
85        }
86    }
87}
88
89impl Index {
90    /// An unnamed implicit `0..n` range index.
91    pub fn range(n: usize) -> Index {
92        Index {
93            name: None,
94            kind: IndexKind::Range(n),
95        }
96    }
97
98    /// An unnamed explicit integer index.
99    pub fn int64(labels: Vec<i64>) -> Index {
100        Index {
101            name: None,
102            kind: IndexKind::Int64(labels),
103        }
104    }
105
106    /// An unnamed datetime index (UTC-ns `labels`, tagged with `tz`).
107    pub fn datetime(labels: Vec<i64>, tz: Tz) -> Index {
108        Index {
109            name: None,
110            kind: IndexKind::Datetime(labels, tz),
111        }
112    }
113
114    /// An unnamed string index.
115    pub fn str(labels: Vec<String>) -> Index {
116        Index {
117            name: None,
118            kind: IndexKind::Str(labels),
119        }
120    }
121
122    /// The label storage / kind.
123    pub fn kind(&self) -> &IndexKind {
124        &self.kind
125    }
126
127    /// The index name, if any (pandas `Index.name`).
128    pub fn name(&self) -> Option<&str> {
129        self.name.as_deref()
130    }
131
132    /// Return this index with its name set (builder; consumes `self`).
133    pub fn with_name(mut self, name: Option<String>) -> Index {
134        self.name = name;
135        self
136    }
137
138    /// Build an unnamed index from a column (for `set_index`): a `Datetime`
139    /// column becomes a `DatetimeIndex`, an `I64` column an `Int64Index`, a `Str`
140    /// column a string index. Float / bool columns are not valid labels.
141    pub fn from_column(col: &Column) -> Result<Index> {
142        Index::from_column_tz(col, Tz::Naive)
143    }
144
145    /// Build an unnamed index from a column, tagging a `Datetime` column with
146    /// `tz` (otherwise the tz is ignored).
147    ///
148    /// An `I64` / `Str` column carrying `volas.NA` is rejected: an int/str index
149    /// has no missing-label representation, so building one would silently turn
150    /// each NA into its physical placeholder (`0` / `""`) — an ordinary label a
151    /// `.loc` lookup can match (C2/C4). Datetime is the one nullable index kind:
152    /// `NaT` is a physical sentinel inside the label vector itself, rendered as
153    /// `NaT` and sorted last.
154    pub fn from_column_tz(col: &Column, tz: Tz) -> Result<Index> {
155        let kind = match col {
156            // Datetime is exempt from the unique-label rule (F34 refinement):
157            // real market data legitimately carries duplicate timestamps (resent
158            // forming bars, multiple ticks per ts, NaT batches) and cumulate /
159            // sort own the dedup semantics. int64/str labels stay strictly unique.
160            Column::Datetime(v) => IndexKind::Datetime(v.to_vec(), tz),
161            Column::I64(v, _) => {
162                require_no_missing_labels(col, "int64")?;
163                require_unique_labels(v, "int64")?;
164                IndexKind::Int64(v.to_vec())
165            }
166            Column::Str(v, _) => {
167                require_no_missing_labels(col, "str")?;
168                let labels = v.to_vec();
169                require_unique_labels(&labels, "str")?;
170                IndexKind::Str(labels)
171            }
172            other => {
173                return Err(VolasError::DType(format!(
174                    "cannot use a {} column as an index (only datetime / int64 / string)",
175                    other.dtype()
176                )))
177            }
178        };
179        Ok(Index { name: None, kind })
180    }
181
182    /// The timezone of a `Datetime` index ([`Tz::Naive`] for every other kind).
183    pub fn tz(&self) -> Tz {
184        match &self.kind {
185            IndexKind::Datetime(_, tz) => *tz,
186            _ => Tz::Naive,
187        }
188    }
189
190    /// Return this index with its timezone set (no-op for a non-datetime index);
191    /// the name is preserved.
192    pub fn with_tz(mut self, tz: Tz) -> Index {
193        if let IndexKind::Datetime(_, cur) = &mut self.kind {
194            *cur = tz;
195        }
196        self
197    }
198
199    /// Number of labels.
200    pub fn len(&self) -> usize {
201        match &self.kind {
202            IndexKind::Range(n) => *n,
203            IndexKind::Int64(v) => v.len(),
204            IndexKind::Datetime(v, _) => v.len(),
205            IndexKind::Str(v) => v.len(),
206        }
207    }
208
209    /// The label at position `i` (for membership tests like `drop`).
210    pub fn label_at(&self, i: usize) -> Label {
211        match &self.kind {
212            IndexKind::Range(_) => Label::I64(i as i64),
213            IndexKind::Int64(v) => Label::I64(v[i]),
214            IndexKind::Datetime(v, _) => Label::I64(v[i]),
215            IndexKind::Str(v) => Label::Str(v[i].clone()),
216        }
217    }
218
219    /// Whether the index is empty.
220    pub fn is_empty(&self) -> bool {
221        self.len() == 0
222    }
223
224    /// Materialize the numeric labels as `i64`. Numeric indexes only — string
225    /// indexes are handled by their own paths (`label_slice` / `append` guard).
226    pub fn to_i64_labels(&self) -> Vec<i64> {
227        self.kind.to_i64_labels()
228    }
229
230    /// A `[start, end)` slice (the name and a datetime tz are preserved).
231    pub fn slice(&self, start: usize, end: usize) -> Index {
232        let kind = match &self.kind {
233            // A range starting at 0 stays an (implicit) range; a non-zero start
234            // cannot be expressed as `Range`, so it materializes the actual labels
235            // (`start..end`) — e.g. `tail(2)` of `0..7` keeps labels `[5, 6]`, not
236            // a reset `[0, 1]`, matching pandas.
237            IndexKind::Range(_) if start == 0 => IndexKind::Range(end),
238            IndexKind::Range(_) => IndexKind::Int64((start as i64..end as i64).collect()),
239            IndexKind::Int64(v) => IndexKind::Int64(v[start..end].to_vec()),
240            IndexKind::Datetime(v, tz) => IndexKind::Datetime(v[start..end].to_vec(), *tz),
241            IndexKind::Str(v) => IndexKind::Str(v[start..end].to_vec()),
242        };
243        Index {
244            name: self.name.clone(),
245            kind,
246        }
247    }
248
249    /// Gather the given positions (the name and a datetime tz are preserved).
250    pub fn take(&self, idx: &[usize]) -> Index {
251        let kind = match &self.kind {
252            IndexKind::Range(_) => IndexKind::Int64(idx.iter().map(|&i| i as i64).collect()),
253            IndexKind::Int64(v) => IndexKind::Int64(idx.iter().map(|&i| v[i]).collect()),
254            IndexKind::Datetime(v, tz) => {
255                IndexKind::Datetime(idx.iter().map(|&i| v[i]).collect(), *tz)
256            }
257            IndexKind::Str(v) => IndexKind::Str(idx.iter().map(|&i| v[i].clone()).collect()),
258        };
259        Index {
260            name: self.name.clone(),
261            kind,
262        }
263    }
264
265    /// Value equality of the **labels** (pandas `.equals` index semantics): a
266    /// `RangeIndex` equals the same integer labels materialized as `Int64`, but an
267    /// integer index never equals a datetime or string index. The index *name* and
268    /// a datetime *tz* (display metadata) are ignored.
269    pub fn label_eq(&self, other: &Index) -> bool {
270        use IndexKind::*;
271        match (&self.kind, &other.kind) {
272            (Range(_) | Int64(_), Range(_) | Int64(_)) => {
273                self.to_i64_labels() == other.to_i64_labels()
274            }
275            (Datetime(a, _), Datetime(b, _)) => a == b,
276            (Str(a), Str(b)) => a == b,
277            _ => false, // different label kinds are never equal
278        }
279    }
280
281    /// The positions that sort the index by label (`sort_index`). Numeric kinds
282    /// sort numerically, a string index lexicographically.
283    pub fn argsort(&self, ascending: bool) -> Vec<usize> {
284        let mut idx: Vec<usize> = (0..self.len()).collect();
285        let cmp_dir = |o: std::cmp::Ordering| if ascending { o } else { o.reverse() };
286        match &self.kind {
287            IndexKind::Str(v) => idx.sort_by(|&a, &b| cmp_dir(v[a].cmp(&v[b]))),
288            // a datetime index sinks a NaT label (i64::MIN) to the end in both
289            // directions — matching sort_values / pandas na_position='last' (V9) —
290            // rather than sorting it as the smallest i64, which put NaT first.
291            IndexKind::Datetime(v, _) => idx.sort_by(|&a, &b| {
292                use std::cmp::Ordering::*;
293                match (v[a] == i64::MIN, v[b] == i64::MIN) {
294                    (true, true) => Equal,
295                    (true, false) => Greater,
296                    (false, true) => Less,
297                    (false, false) => cmp_dir(v[a].cmp(&v[b])),
298                }
299            }),
300            _ => {
301                let labels = self.to_i64_labels();
302                idx.sort_by(|&a, &b| cmp_dir(labels[a].cmp(&labels[b])));
303            }
304        }
305        idx
306    }
307
308    /// Materialize the labels as a [`Column`] (for `reset_index`).
309    pub fn to_column(&self) -> Column {
310        match &self.kind {
311            IndexKind::Range(n) => Column::i64((0..*n as i64).collect()),
312            IndexKind::Int64(v) => Column::i64(v.clone()),
313            IndexKind::Datetime(v, _) => Column::datetime(v.clone()),
314            IndexKind::Str(v) => Column::str(v.clone()),
315        }
316    }
317
318    /// Concatenate two indexes (extending labels), keeping the left index's name.
319    /// Same-kind indexes preserve their kind; mixing numeric kinds yields
320    /// `Int64`; mixing a string index with a numeric one is an error.
321    pub fn append(&self, other: &Index) -> Result<Index> {
322        use IndexKind::*;
323        let kind = match (&self.kind, &other.kind) {
324            (Range(a), Range(b)) => Range(a + b),
325            (Datetime(a, ta), Datetime(b, _)) => Datetime([a.as_slice(), b].concat(), *ta),
326            (Str(a), Str(b)) => Str([a.as_slice(), b].concat()),
327            (Str(_), _) | (_, Str(_)) => {
328                return Err(VolasError::Shape(
329                    "cannot append a string index to a non-string index".into(),
330                ))
331            }
332            // remaining: numeric mixes (Range / Int64 / Datetime) -> Int64 labels
333            (a, b) => Int64([a.to_i64_labels(), b.to_i64_labels()].concat()),
334        };
335        Ok(Index {
336            name: self.name.clone(),
337            kind,
338        })
339    }
340
341    /// Extend in place by the labels of `other` — the amortized-O(1) counterpart
342    /// of [`append`](Self::append), used by the live single-bar hot path; the
343    /// growing index keeps its own name (so appending an unnamed bar to a named
344    /// index does not drop the name). Same-kind indexes grow their buffer; a
345    /// numeric-kind mix collapses to `Int64`; mixing a string index with a
346    /// numeric one is an error.
347    pub fn extend(&mut self, other: &Index) -> Result<()> {
348        use IndexKind::*;
349        match (&mut self.kind, &other.kind) {
350            (Range(a), Range(b)) => *a += b,
351            (Datetime(a, _), Datetime(b, _)) => a.extend_from_slice(b),
352            (Int64(a), Int64(b)) => a.extend_from_slice(b),
353            (Str(a), Str(b)) => a.extend(b.iter().cloned()),
354            (Str(_), _) | (_, Str(_)) => {
355                return Err(VolasError::Shape(
356                    "cannot append a string index to a non-string index".into(),
357                ))
358            }
359            // An implicit positional (Range) index appended to explicit Int64 labels
360            // continues the sequence (last + 1, ...), mirroring Range+Range's auto-
361            // increment — so a sliced / compacted frame (whose Range collapsed to
362            // absolute Int64 labels) keeps assigning monotonic labels to appended bars
363            // instead of restarting at the bar's own 0. The bare-RangeIndex bar is the
364            // "no explicit label, take the next row position" signal; honoring it keeps
365            // a windowed (rolling-compacted) RangeIndex frame's labels absolute and
366            // gap-free, matching pandas' last-M view of a RangeIndex.
367            (Int64(a), Range(bn)) => {
368                let next = a.last().map(|l| l + 1).unwrap_or(0);
369                a.extend((0..*bn as i64).map(|k| next + k));
370            }
371            // numeric-kind mix (Range / Int64 / Datetime) -> Int64 labels
372            (slot, b) => {
373                let mut labels = slot.to_i64_labels();
374                labels.extend(b.to_i64_labels());
375                *slot = Int64(labels);
376            }
377        }
378        Ok(())
379    }
380
381    /// Position of the first label exactly equal to `label`. Returns `None` if
382    /// the label's kind does not match the index's kind.
383    pub fn position_of(&self, label: &Label) -> Option<usize> {
384        match (&self.kind, label) {
385            (IndexKind::Range(n), Label::I64(v)) => {
386                if *v >= 0 && (*v as usize) < *n {
387                    Some(*v as usize)
388                } else {
389                    None
390                }
391            }
392            (IndexKind::Int64(vs), Label::I64(v)) => vs.iter().position(|x| x == v),
393            (IndexKind::Datetime(vs, _), Label::I64(v)) => vs.iter().position(|x| x == v),
394            (IndexKind::Str(vs), Label::Str(s)) => vs.iter().position(|x| x == s),
395            _ => None,
396        }
397    }
398
399    /// `[start, end)` positions covering the inclusive label range `[lo, hi]`
400    /// (ascending labels; pandas `.loc` slice semantics). Either bound may be
401    /// `None` for open-ended. Numeric indexes compare numerically; a string
402    /// index compares lexicographically.
403    pub fn label_slice(&self, lo: Option<&Label>, hi: Option<&Label>) -> (usize, usize) {
404        match &self.kind {
405            IndexKind::Str(labels) => {
406                let start = lo.and_then(Label::as_str).map_or(0, |lo| {
407                    labels
408                        .iter()
409                        .position(|x| x.as_str() >= lo)
410                        .unwrap_or(labels.len())
411                });
412                let end = hi.and_then(Label::as_str).map_or(labels.len(), |hi| {
413                    labels
414                        .iter()
415                        .rposition(|x| x.as_str() <= hi)
416                        .map_or(0, |p| p + 1)
417                });
418                (start, end.max(start))
419            }
420            _ => {
421                let labels = self.to_i64_labels();
422                let start = lo.and_then(Label::as_i64).map_or(0, |lo| {
423                    labels.iter().position(|&x| x >= lo).unwrap_or(labels.len())
424                });
425                let end = hi.and_then(Label::as_i64).map_or(labels.len(), |hi| {
426                    labels.iter().rposition(|&x| x <= hi).map_or(0, |p| p + 1)
427                });
428                (start, end.max(start))
429            }
430        }
431    }
432}
433
434/// Reject building an index from a column that carries `volas.NA`: with no
435/// missing-label representation for `int64` / `str`, the NA's physical
436/// placeholder (`0` / `""`) would become an ordinary, lookup-matchable label.
437fn require_no_missing_labels(col: &Column, kind: &str) -> Result<()> {
438    if col.null_count() > 0 {
439        return Err(VolasError::Value(format!(
440            "cannot use a {kind} column containing volas.NA as an index (a missing \
441             label has no {kind} representation); drop or fill the NA rows first"
442        )));
443    }
444    Ok(())
445}
446
447/// Label access assumes unique labels (F34, decision 1B): a duplicate-label
448/// index makes `.loc[label]` ill-defined (one row or many?), so it is rejected
449/// at creation — the same creation-time guard as the NA-label rule above.
450fn require_unique_labels<T: std::hash::Hash + Eq>(labels: &[T], kind: &str) -> Result<()> {
451    let mut seen = std::collections::HashSet::with_capacity(labels.len());
452    for l in labels {
453        if !seen.insert(l) {
454            return Err(VolasError::Value(format!(
455                "cannot use a {kind} column with duplicate labels as an index \
456                 (label access assumes unique labels)"
457            )));
458        }
459    }
460    Ok(())
461}
462
463#[cfg(test)]
464mod tests {
465    use super::*;
466
467    #[test]
468    fn from_datetime_and_int_columns() {
469        assert_eq!(
470            Index::from_column(&Column::datetime(vec![5, 6])).unwrap(),
471            Index::datetime(vec![5, 6], Tz::Naive)
472        );
473        assert_eq!(
474            Index::from_column(&Column::i64(vec![1, 2])).unwrap(),
475            Index::int64(vec![1, 2])
476        );
477    }
478
479    #[test]
480    fn from_unsupported_column_errors() {
481        // float / bool columns are not valid labels (string is — see below)
482        assert!(Index::from_column(&Column::f64(vec![1.0])).is_err());
483        assert!(Index::from_column(&Column::bool(vec![true])).is_err());
484    }
485
486    #[test]
487    fn is_empty_labels_and_position_of() {
488        assert!(Index::range(0).is_empty());
489        assert!(!Index::range(3).is_empty());
490
491        assert_eq!(Index::range(3).to_i64_labels(), vec![0, 1, 2]);
492        assert_eq!(Index::int64(vec![5, 6]).to_i64_labels(), vec![5, 6]);
493        assert_eq!(
494            Index::datetime(vec![10, 20], Tz::Utc).to_i64_labels(),
495            vec![10, 20]
496        );
497
498        let i64 = Label::I64;
499        assert_eq!(Index::range(5).position_of(&i64(3)), Some(3));
500        assert_eq!(Index::range(5).position_of(&i64(9)), None);
501        assert_eq!(Index::range(5).position_of(&i64(-1)), None);
502        assert_eq!(Index::int64(vec![10, 20, 30]).position_of(&i64(20)), Some(1));
503        assert_eq!(Index::int64(vec![10, 20]).position_of(&i64(99)), None);
504        assert_eq!(
505            Index::datetime(vec![100, 200], Tz::Utc).position_of(&i64(200)),
506            Some(1)
507        );
508
509        // take() on an Int64 index gathers the labels at those positions
510        assert_eq!(
511            Index::int64(vec![10, 20, 30]).take(&[2, 0]),
512            Index::int64(vec![30, 10])
513        );
514    }
515
516    fn str_index(labels: &[&str]) -> Index {
517        Index::str(labels.iter().map(|s| s.to_string()).collect())
518    }
519
520    #[test]
521    fn string_index_construction_and_ops() {
522        // a string column becomes a string index
523        let ix = Index::from_column(&Column::str(vec!["a".into(), "b".into()])).unwrap();
524        assert_eq!(ix, str_index(&["a", "b"]));
525
526        let ix = str_index(&["a", "b", "c", "d"]);
527        assert_eq!(ix.len(), 4);
528        assert!(!ix.is_empty());
529        assert_eq!(ix.label_at(2), Label::Str("c".into()));
530        assert_eq!(ix.slice(1, 3), str_index(&["b", "c"]));
531        assert_eq!(ix.take(&[3, 0]), str_index(&["d", "a"]));
532    }
533
534    #[test]
535    fn string_index_lookup_and_slice() {
536        let ix = str_index(&["aa", "bb", "cc", "dd"]);
537        // exact lookup, kind-matched
538        assert_eq!(ix.position_of(&Label::Str("cc".into())), Some(2));
539        assert_eq!(ix.position_of(&Label::Str("zz".into())), None);
540        // a numeric label against a string index never matches
541        assert_eq!(ix.position_of(&Label::I64(1)), None);
542        // lexicographic slice [bb, cc]
543        let lo = Label::Str("bb".into());
544        let hi = Label::Str("cc".into());
545        assert_eq!(ix.label_slice(Some(&lo), Some(&hi)), (1, 3));
546        // open-ended upper bound: [cc, end)
547        assert_eq!(ix.label_slice(Some(&lo), None), (1, 4));
548    }
549
550    #[test]
551    fn string_index_append_rules() {
552        let a = str_index(&["x", "y"]);
553        let b = str_index(&["z"]);
554        assert_eq!(a.append(&b).unwrap(), str_index(&["x", "y", "z"]));
555        // mixing a string index with a numeric one is an error
556        assert!(a.append(&Index::range(2)).is_err());
557        assert!(Index::range(2).append(&a).is_err());
558    }
559
560    #[test]
561    fn extend_grows_in_place_per_kind() {
562        // same-kind grows the buffer in place (the live append hot path)
563        let mut r = Index::range(3);
564        r.extend(&Index::range(2)).unwrap();
565        assert_eq!(r, Index::range(5));
566
567        let mut d = Index::datetime(vec![1, 2], Tz::Utc);
568        d.extend(&Index::datetime(vec![3], Tz::Utc)).unwrap();
569        assert_eq!(d, Index::datetime(vec![1, 2, 3], Tz::Utc));
570
571        let mut s = str_index(&["a", "b"]);
572        s.extend(&str_index(&["c"])).unwrap();
573        assert_eq!(s, str_index(&["a", "b", "c"]));
574
575        // a numeric-kind mix collapses to Int64 (matches `append`)
576        let mut m = Index::range(2);
577        m.extend(&Index::int64(vec![5, 6])).unwrap();
578        assert_eq!(m, Index::int64(vec![0, 1, 5, 6]));
579
580        // an implicit Range appended to explicit Int64 labels CONTINUES the sequence
581        // (last + 1, ...) rather than restarting at the Range's own 0 — the
582        // sliced/compacted-frame + bare-bar append path (windowed rolling buffers).
583        let mut c = Index::int64(vec![44, 45, 46]);
584        c.extend(&Index::range(2)).unwrap();
585        assert_eq!(c, Index::int64(vec![44, 45, 46, 47, 48]));
586        // continuing from an empty Int64 starts at 0.
587        let mut e = Index::int64(vec![]);
588        e.extend(&Index::range(2)).unwrap();
589        assert_eq!(e, Index::int64(vec![0, 1]));
590
591        // mixing string with numeric is an error, either way
592        assert!(str_index(&["x"]).extend(&Index::range(1)).is_err());
593        assert!(Index::range(1).extend(&str_index(&["x"])).is_err());
594    }
595
596    #[test]
597    fn name_set_and_propagates_through_ops() {
598        let ix = Index::datetime(vec![1, 2, 3], Tz::Utc).with_name(Some("date".into()));
599        assert_eq!(ix.name(), Some("date"));
600        // an unnamed index reports None
601        assert_eq!(Index::range(3).name(), None);
602        // the name rides through identity-preserving ops
603        assert_eq!(ix.slice(0, 2).name(), Some("date"));
604        assert_eq!(ix.take(&[2, 0]).name(), Some("date"));
605        assert_eq!(ix.clone().with_tz(Tz::Offset(28800)).name(), Some("date"));
606        // append / extend keep the left (growing) index's name
607        assert_eq!(
608            ix.append(&Index::datetime(vec![4], Tz::Utc)).unwrap().name(),
609            Some("date")
610        );
611        let mut g = ix.clone();
612        g.extend(&Index::datetime(vec![4], Tz::Utc)).unwrap();
613        assert_eq!(g.name(), Some("date"));
614        // with_name(None) clears it
615        assert_eq!(ix.with_name(None).name(), None);
616    }
617
618    #[test]
619    fn label_accessors_and_numeric_label_at() {
620        // accessors return None on the other variant
621        assert_eq!(Label::I64(5).as_i64(), Some(5));
622        assert_eq!(Label::I64(5).as_str(), None);
623        assert_eq!(Label::Str("x".into()).as_str(), Some("x"));
624        assert_eq!(Label::Str("x".into()).as_i64(), None);
625        // label_at over the numeric index kinds
626        assert_eq!(Index::range(3).label_at(2), Label::I64(2));
627        assert_eq!(Index::int64(vec![10, 20]).label_at(1), Label::I64(20));
628        assert_eq!(
629            Index::datetime(vec![100, 200], Tz::Utc).label_at(0),
630            Label::I64(100)
631        );
632    }
633
634    #[test]
635    fn index_kind_branch_coverage() {
636        // tz() / with_tz() are no-ops on a non-datetime index.
637        assert_eq!(Index::range(3).tz(), Tz::Naive);
638        assert!(matches!(
639            Index::range(3).with_tz(Tz::Utc).kind,
640            IndexKind::Range(3)
641        ));
642        // slice over the non-range kinds.
643        assert_eq!(
644            Index::int64(vec![1, 2, 3]).slice(0, 2),
645            Index::int64(vec![1, 2])
646        );
647        assert!(matches!(
648            Index::datetime(vec![1, 2], Tz::Utc).slice(0, 1).kind,
649            IndexKind::Datetime(_, _)
650        ));
651        assert_eq!(
652            Index::str(vec!["a".into(), "b".into()]).slice(1, 2),
653            Index::str(vec!["b".into()])
654        );
655        // argsort lexicographically over a string index.
656        assert_eq!(
657            Index::str(vec!["b".into(), "a".into()]).argsort(true),
658            vec![1, 0]
659        );
660        // to_column over every kind.
661        assert_eq!(Index::range(2).to_column().len(), 2);
662        assert_eq!(Index::datetime(vec![5], Tz::Utc).to_column().len(), 1);
663        assert_eq!(Index::str(vec!["x".into()]).to_column().len(), 1);
664        // append: same-kind datetime / string, and a numeric mix -> Int64.
665        assert!(matches!(
666            Index::datetime(vec![1], Tz::Utc)
667                .append(&Index::datetime(vec![2], Tz::Utc))
668                .unwrap()
669                .kind,
670            IndexKind::Datetime(_, _)
671        ));
672        assert!(matches!(
673            Index::str(vec!["a".into()])
674                .append(&Index::str(vec!["b".into()]))
675                .unwrap()
676                .kind,
677            IndexKind::Str(_)
678        ));
679        assert!(matches!(
680            Index::range(2).append(&Index::range(3)).unwrap().kind,
681            IndexKind::Range(5)
682        ));
683        assert!(matches!(
684            Index::range(2).append(&Index::int64(vec![5])).unwrap().kind,
685            IndexKind::Int64(_)
686        ));
687        // mixing a string index with a numeric one is an error.
688        assert!(Index::str(vec!["a".into()])
689            .append(&Index::range(1))
690            .is_err());
691    }
692
693    #[test]
694    fn label_eq_value_semantics() {
695        // a RangeIndex equals the same integer labels materialized as Int64
696        assert!(Index::range(3).label_eq(&Index::int64(vec![0, 1, 2])));
697        assert!(!Index::range(3).label_eq(&Index::int64(vec![0, 1, 9])));
698        // datetime and string indexes compare by value
699        assert!(Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 2], Tz::Utc)));
700        assert!(!Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 9], Tz::Utc)));
701        assert!(Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["a".into()])));
702        assert!(!Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["b".into()])));
703        // different label kinds are never equal
704        assert!(!Index::range(2).label_eq(&Index::datetime(vec![0, 1], Tz::Utc)));
705        assert!(!Index::str(vec!["a".into()]).label_eq(&Index::int64(vec![0])));
706    }
707}