1use crate::column::Column;
4use crate::error::{Result, VolasError};
5use crate::tz::Tz;
6
7#[derive(Clone, Debug, PartialEq)]
20pub struct Index {
21 pub name: Option<String>,
23 pub kind: IndexKind,
25}
26
27#[derive(Clone, Debug, PartialEq)]
33pub enum IndexKind {
34 Range(usize),
36 Int64(Vec<i64>),
38 Datetime(Vec<i64>, Tz),
41 Str(Vec<String>),
43}
44
45impl IndexKind {
46 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"), }
55 }
56}
57
58#[derive(Clone, Debug, PartialEq)]
62pub enum Label {
63 I64(i64),
65 Str(String),
67}
68
69impl Label {
70 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 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 pub fn range(n: usize) -> Index {
92 Index {
93 name: None,
94 kind: IndexKind::Range(n),
95 }
96 }
97
98 pub fn int64(labels: Vec<i64>) -> Index {
100 Index {
101 name: None,
102 kind: IndexKind::Int64(labels),
103 }
104 }
105
106 pub fn datetime(labels: Vec<i64>, tz: Tz) -> Index {
108 Index {
109 name: None,
110 kind: IndexKind::Datetime(labels, tz),
111 }
112 }
113
114 pub fn str(labels: Vec<String>) -> Index {
116 Index {
117 name: None,
118 kind: IndexKind::Str(labels),
119 }
120 }
121
122 pub fn kind(&self) -> &IndexKind {
124 &self.kind
125 }
126
127 pub fn name(&self) -> Option<&str> {
129 self.name.as_deref()
130 }
131
132 pub fn with_name(mut self, name: Option<String>) -> Index {
134 self.name = name;
135 self
136 }
137
138 pub fn from_column(col: &Column) -> Result<Index> {
142 Index::from_column_tz(col, Tz::Naive)
143 }
144
145 pub fn from_column_tz(col: &Column, tz: Tz) -> Result<Index> {
155 let kind = match col {
156 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 pub fn tz(&self) -> Tz {
184 match &self.kind {
185 IndexKind::Datetime(_, tz) => *tz,
186 _ => Tz::Naive,
187 }
188 }
189
190 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 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 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 pub fn is_empty(&self) -> bool {
221 self.len() == 0
222 }
223
224 pub fn to_i64_labels(&self) -> Vec<i64> {
227 self.kind.to_i64_labels()
228 }
229
230 pub fn slice(&self, start: usize, end: usize) -> Index {
232 let kind = match &self.kind {
233 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 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 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, }
279 }
280
281 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 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 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 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 (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 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 (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 (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 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 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
434fn 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
447fn 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 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 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 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 assert_eq!(ix.position_of(&Label::Str("cc".into())), Some(2));
539 assert_eq!(ix.position_of(&Label::Str("zz".into())), None);
540 assert_eq!(ix.position_of(&Label::I64(1)), None);
542 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 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 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 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 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 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 let mut e = Index::int64(vec![]);
588 e.extend(&Index::range(2)).unwrap();
589 assert_eq!(e, Index::int64(vec![0, 1]));
590
591 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 assert_eq!(Index::range(3).name(), None);
602 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 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 assert_eq!(ix.with_name(None).name(), None);
616 }
617
618 #[test]
619 fn label_accessors_and_numeric_label_at() {
620 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 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 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 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 assert_eq!(
657 Index::str(vec!["b".into(), "a".into()]).argsort(true),
658 vec![1, 0]
659 );
660 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 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 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 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 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 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}