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 (slot, b) => {
361 let mut labels = slot.to_i64_labels();
362 labels.extend(b.to_i64_labels());
363 *slot = Int64(labels);
364 }
365 }
366 Ok(())
367 }
368
369 pub fn position_of(&self, label: &Label) -> Option<usize> {
372 match (&self.kind, label) {
373 (IndexKind::Range(n), Label::I64(v)) => {
374 if *v >= 0 && (*v as usize) < *n {
375 Some(*v as usize)
376 } else {
377 None
378 }
379 }
380 (IndexKind::Int64(vs), Label::I64(v)) => vs.iter().position(|x| x == v),
381 (IndexKind::Datetime(vs, _), Label::I64(v)) => vs.iter().position(|x| x == v),
382 (IndexKind::Str(vs), Label::Str(s)) => vs.iter().position(|x| x == s),
383 _ => None,
384 }
385 }
386
387 pub fn label_slice(&self, lo: Option<&Label>, hi: Option<&Label>) -> (usize, usize) {
392 match &self.kind {
393 IndexKind::Str(labels) => {
394 let start = lo.and_then(Label::as_str).map_or(0, |lo| {
395 labels
396 .iter()
397 .position(|x| x.as_str() >= lo)
398 .unwrap_or(labels.len())
399 });
400 let end = hi.and_then(Label::as_str).map_or(labels.len(), |hi| {
401 labels
402 .iter()
403 .rposition(|x| x.as_str() <= hi)
404 .map_or(0, |p| p + 1)
405 });
406 (start, end.max(start))
407 }
408 _ => {
409 let labels = self.to_i64_labels();
410 let start = lo.and_then(Label::as_i64).map_or(0, |lo| {
411 labels.iter().position(|&x| x >= lo).unwrap_or(labels.len())
412 });
413 let end = hi.and_then(Label::as_i64).map_or(labels.len(), |hi| {
414 labels.iter().rposition(|&x| x <= hi).map_or(0, |p| p + 1)
415 });
416 (start, end.max(start))
417 }
418 }
419 }
420}
421
422fn require_no_missing_labels(col: &Column, kind: &str) -> Result<()> {
426 if col.null_count() > 0 {
427 return Err(VolasError::Value(format!(
428 "cannot use a {kind} column containing volas.NA as an index (a missing \
429 label has no {kind} representation); drop or fill the NA rows first"
430 )));
431 }
432 Ok(())
433}
434
435fn require_unique_labels<T: std::hash::Hash + Eq>(labels: &[T], kind: &str) -> Result<()> {
439 let mut seen = std::collections::HashSet::with_capacity(labels.len());
440 for l in labels {
441 if !seen.insert(l) {
442 return Err(VolasError::Value(format!(
443 "cannot use a {kind} column with duplicate labels as an index \
444 (label access assumes unique labels)"
445 )));
446 }
447 }
448 Ok(())
449}
450
451#[cfg(test)]
452mod tests {
453 use super::*;
454
455 #[test]
456 fn from_datetime_and_int_columns() {
457 assert_eq!(
458 Index::from_column(&Column::datetime(vec![5, 6])).unwrap(),
459 Index::datetime(vec![5, 6], Tz::Naive)
460 );
461 assert_eq!(
462 Index::from_column(&Column::i64(vec![1, 2])).unwrap(),
463 Index::int64(vec![1, 2])
464 );
465 }
466
467 #[test]
468 fn from_unsupported_column_errors() {
469 assert!(Index::from_column(&Column::f64(vec![1.0])).is_err());
471 assert!(Index::from_column(&Column::bool(vec![true])).is_err());
472 }
473
474 #[test]
475 fn is_empty_labels_and_position_of() {
476 assert!(Index::range(0).is_empty());
477 assert!(!Index::range(3).is_empty());
478
479 assert_eq!(Index::range(3).to_i64_labels(), vec![0, 1, 2]);
480 assert_eq!(Index::int64(vec![5, 6]).to_i64_labels(), vec![5, 6]);
481 assert_eq!(
482 Index::datetime(vec![10, 20], Tz::Utc).to_i64_labels(),
483 vec![10, 20]
484 );
485
486 let i64 = Label::I64;
487 assert_eq!(Index::range(5).position_of(&i64(3)), Some(3));
488 assert_eq!(Index::range(5).position_of(&i64(9)), None);
489 assert_eq!(Index::range(5).position_of(&i64(-1)), None);
490 assert_eq!(Index::int64(vec![10, 20, 30]).position_of(&i64(20)), Some(1));
491 assert_eq!(Index::int64(vec![10, 20]).position_of(&i64(99)), None);
492 assert_eq!(
493 Index::datetime(vec![100, 200], Tz::Utc).position_of(&i64(200)),
494 Some(1)
495 );
496
497 assert_eq!(
499 Index::int64(vec![10, 20, 30]).take(&[2, 0]),
500 Index::int64(vec![30, 10])
501 );
502 }
503
504 fn str_index(labels: &[&str]) -> Index {
505 Index::str(labels.iter().map(|s| s.to_string()).collect())
506 }
507
508 #[test]
509 fn string_index_construction_and_ops() {
510 let ix = Index::from_column(&Column::str(vec!["a".into(), "b".into()])).unwrap();
512 assert_eq!(ix, str_index(&["a", "b"]));
513
514 let ix = str_index(&["a", "b", "c", "d"]);
515 assert_eq!(ix.len(), 4);
516 assert!(!ix.is_empty());
517 assert_eq!(ix.label_at(2), Label::Str("c".into()));
518 assert_eq!(ix.slice(1, 3), str_index(&["b", "c"]));
519 assert_eq!(ix.take(&[3, 0]), str_index(&["d", "a"]));
520 }
521
522 #[test]
523 fn string_index_lookup_and_slice() {
524 let ix = str_index(&["aa", "bb", "cc", "dd"]);
525 assert_eq!(ix.position_of(&Label::Str("cc".into())), Some(2));
527 assert_eq!(ix.position_of(&Label::Str("zz".into())), None);
528 assert_eq!(ix.position_of(&Label::I64(1)), None);
530 let lo = Label::Str("bb".into());
532 let hi = Label::Str("cc".into());
533 assert_eq!(ix.label_slice(Some(&lo), Some(&hi)), (1, 3));
534 assert_eq!(ix.label_slice(Some(&lo), None), (1, 4));
536 }
537
538 #[test]
539 fn string_index_append_rules() {
540 let a = str_index(&["x", "y"]);
541 let b = str_index(&["z"]);
542 assert_eq!(a.append(&b).unwrap(), str_index(&["x", "y", "z"]));
543 assert!(a.append(&Index::range(2)).is_err());
545 assert!(Index::range(2).append(&a).is_err());
546 }
547
548 #[test]
549 fn extend_grows_in_place_per_kind() {
550 let mut r = Index::range(3);
552 r.extend(&Index::range(2)).unwrap();
553 assert_eq!(r, Index::range(5));
554
555 let mut d = Index::datetime(vec![1, 2], Tz::Utc);
556 d.extend(&Index::datetime(vec![3], Tz::Utc)).unwrap();
557 assert_eq!(d, Index::datetime(vec![1, 2, 3], Tz::Utc));
558
559 let mut s = str_index(&["a", "b"]);
560 s.extend(&str_index(&["c"])).unwrap();
561 assert_eq!(s, str_index(&["a", "b", "c"]));
562
563 let mut m = Index::range(2);
565 m.extend(&Index::int64(vec![5, 6])).unwrap();
566 assert_eq!(m, Index::int64(vec![0, 1, 5, 6]));
567
568 assert!(str_index(&["x"]).extend(&Index::range(1)).is_err());
570 assert!(Index::range(1).extend(&str_index(&["x"])).is_err());
571 }
572
573 #[test]
574 fn name_set_and_propagates_through_ops() {
575 let ix = Index::datetime(vec![1, 2, 3], Tz::Utc).with_name(Some("date".into()));
576 assert_eq!(ix.name(), Some("date"));
577 assert_eq!(Index::range(3).name(), None);
579 assert_eq!(ix.slice(0, 2).name(), Some("date"));
581 assert_eq!(ix.take(&[2, 0]).name(), Some("date"));
582 assert_eq!(ix.clone().with_tz(Tz::Offset(28800)).name(), Some("date"));
583 assert_eq!(
585 ix.append(&Index::datetime(vec![4], Tz::Utc)).unwrap().name(),
586 Some("date")
587 );
588 let mut g = ix.clone();
589 g.extend(&Index::datetime(vec![4], Tz::Utc)).unwrap();
590 assert_eq!(g.name(), Some("date"));
591 assert_eq!(ix.with_name(None).name(), None);
593 }
594
595 #[test]
596 fn label_accessors_and_numeric_label_at() {
597 assert_eq!(Label::I64(5).as_i64(), Some(5));
599 assert_eq!(Label::I64(5).as_str(), None);
600 assert_eq!(Label::Str("x".into()).as_str(), Some("x"));
601 assert_eq!(Label::Str("x".into()).as_i64(), None);
602 assert_eq!(Index::range(3).label_at(2), Label::I64(2));
604 assert_eq!(Index::int64(vec![10, 20]).label_at(1), Label::I64(20));
605 assert_eq!(
606 Index::datetime(vec![100, 200], Tz::Utc).label_at(0),
607 Label::I64(100)
608 );
609 }
610
611 #[test]
612 fn index_kind_branch_coverage() {
613 assert_eq!(Index::range(3).tz(), Tz::Naive);
615 assert!(matches!(
616 Index::range(3).with_tz(Tz::Utc).kind,
617 IndexKind::Range(3)
618 ));
619 assert_eq!(
621 Index::int64(vec![1, 2, 3]).slice(0, 2),
622 Index::int64(vec![1, 2])
623 );
624 assert!(matches!(
625 Index::datetime(vec![1, 2], Tz::Utc).slice(0, 1).kind,
626 IndexKind::Datetime(_, _)
627 ));
628 assert_eq!(
629 Index::str(vec!["a".into(), "b".into()]).slice(1, 2),
630 Index::str(vec!["b".into()])
631 );
632 assert_eq!(
634 Index::str(vec!["b".into(), "a".into()]).argsort(true),
635 vec![1, 0]
636 );
637 assert_eq!(Index::range(2).to_column().len(), 2);
639 assert_eq!(Index::datetime(vec![5], Tz::Utc).to_column().len(), 1);
640 assert_eq!(Index::str(vec!["x".into()]).to_column().len(), 1);
641 assert!(matches!(
643 Index::datetime(vec![1], Tz::Utc)
644 .append(&Index::datetime(vec![2], Tz::Utc))
645 .unwrap()
646 .kind,
647 IndexKind::Datetime(_, _)
648 ));
649 assert!(matches!(
650 Index::str(vec!["a".into()])
651 .append(&Index::str(vec!["b".into()]))
652 .unwrap()
653 .kind,
654 IndexKind::Str(_)
655 ));
656 assert!(matches!(
657 Index::range(2).append(&Index::range(3)).unwrap().kind,
658 IndexKind::Range(5)
659 ));
660 assert!(matches!(
661 Index::range(2).append(&Index::int64(vec![5])).unwrap().kind,
662 IndexKind::Int64(_)
663 ));
664 assert!(Index::str(vec!["a".into()])
666 .append(&Index::range(1))
667 .is_err());
668 }
669
670 #[test]
671 fn label_eq_value_semantics() {
672 assert!(Index::range(3).label_eq(&Index::int64(vec![0, 1, 2])));
674 assert!(!Index::range(3).label_eq(&Index::int64(vec![0, 1, 9])));
675 assert!(Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 2], Tz::Utc)));
677 assert!(!Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 9], Tz::Utc)));
678 assert!(Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["a".into()])));
679 assert!(!Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["b".into()])));
680 assert!(!Index::range(2).label_eq(&Index::datetime(vec![0, 1], Tz::Utc)));
682 assert!(!Index::str(vec!["a".into()]).label_eq(&Index::int64(vec![0])));
683 }
684}