use crate::column::Column;
use crate::error::{Result, VolasError};
use crate::tz::Tz;
#[derive(Clone, Debug, PartialEq)]
pub struct Index {
pub name: Option<String>,
pub kind: IndexKind,
}
#[derive(Clone, Debug, PartialEq)]
pub enum IndexKind {
Range(usize),
Int64(Vec<i64>),
Datetime(Vec<i64>, Tz),
Str(Vec<String>),
}
impl IndexKind {
fn to_i64_labels(&self) -> Vec<i64> {
match self {
IndexKind::Range(n) => (0..*n as i64).collect(),
IndexKind::Int64(v) => v.clone(),
IndexKind::Datetime(v, _) => v.clone(),
IndexKind::Str(_) => unreachable!("string indexes have no i64 labels"), }
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum Label {
I64(i64),
Str(String),
}
impl Label {
pub fn as_i64(&self) -> Option<i64> {
if let Label::I64(v) = self {
Some(*v)
} else {
None
}
}
pub fn as_str(&self) -> Option<&str> {
if let Label::Str(s) = self {
Some(s.as_str())
} else {
None
}
}
}
impl Index {
pub fn range(n: usize) -> Index {
Index {
name: None,
kind: IndexKind::Range(n),
}
}
pub fn int64(labels: Vec<i64>) -> Index {
Index {
name: None,
kind: IndexKind::Int64(labels),
}
}
pub fn datetime(labels: Vec<i64>, tz: Tz) -> Index {
Index {
name: None,
kind: IndexKind::Datetime(labels, tz),
}
}
pub fn str(labels: Vec<String>) -> Index {
Index {
name: None,
kind: IndexKind::Str(labels),
}
}
pub fn kind(&self) -> &IndexKind {
&self.kind
}
pub fn name(&self) -> Option<&str> {
self.name.as_deref()
}
pub fn with_name(mut self, name: Option<String>) -> Index {
self.name = name;
self
}
pub fn from_column(col: &Column) -> Result<Index> {
Index::from_column_tz(col, Tz::Naive)
}
pub fn from_column_tz(col: &Column, tz: Tz) -> Result<Index> {
let kind = match col {
Column::Datetime(v) => IndexKind::Datetime(v.to_vec(), tz),
Column::I64(v, _) => {
require_no_missing_labels(col, "int64")?;
require_unique_labels(v, "int64")?;
IndexKind::Int64(v.to_vec())
}
Column::Str(v, _) => {
require_no_missing_labels(col, "str")?;
require_unique_labels(v, "str")?;
IndexKind::Str(v.to_vec())
}
other => {
return Err(VolasError::DType(format!(
"cannot use a {} column as an index (only datetime / int64 / string)",
other.dtype()
)))
}
};
Ok(Index { name: None, kind })
}
pub fn tz(&self) -> Tz {
match &self.kind {
IndexKind::Datetime(_, tz) => *tz,
_ => Tz::Naive,
}
}
pub fn with_tz(mut self, tz: Tz) -> Index {
if let IndexKind::Datetime(_, cur) = &mut self.kind {
*cur = tz;
}
self
}
pub fn len(&self) -> usize {
match &self.kind {
IndexKind::Range(n) => *n,
IndexKind::Int64(v) => v.len(),
IndexKind::Datetime(v, _) => v.len(),
IndexKind::Str(v) => v.len(),
}
}
pub fn label_at(&self, i: usize) -> Label {
match &self.kind {
IndexKind::Range(_) => Label::I64(i as i64),
IndexKind::Int64(v) => Label::I64(v[i]),
IndexKind::Datetime(v, _) => Label::I64(v[i]),
IndexKind::Str(v) => Label::Str(v[i].clone()),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn to_i64_labels(&self) -> Vec<i64> {
self.kind.to_i64_labels()
}
pub fn slice(&self, start: usize, end: usize) -> Index {
let kind = match &self.kind {
IndexKind::Range(_) if start == 0 => IndexKind::Range(end),
IndexKind::Range(_) => IndexKind::Int64((start as i64..end as i64).collect()),
IndexKind::Int64(v) => IndexKind::Int64(v[start..end].to_vec()),
IndexKind::Datetime(v, tz) => IndexKind::Datetime(v[start..end].to_vec(), *tz),
IndexKind::Str(v) => IndexKind::Str(v[start..end].to_vec()),
};
Index {
name: self.name.clone(),
kind,
}
}
pub fn take(&self, idx: &[usize]) -> Index {
let kind = match &self.kind {
IndexKind::Range(_) => IndexKind::Int64(idx.iter().map(|&i| i as i64).collect()),
IndexKind::Int64(v) => IndexKind::Int64(idx.iter().map(|&i| v[i]).collect()),
IndexKind::Datetime(v, tz) => {
IndexKind::Datetime(idx.iter().map(|&i| v[i]).collect(), *tz)
}
IndexKind::Str(v) => IndexKind::Str(idx.iter().map(|&i| v[i].clone()).collect()),
};
Index {
name: self.name.clone(),
kind,
}
}
pub fn label_eq(&self, other: &Index) -> bool {
use IndexKind::*;
match (&self.kind, &other.kind) {
(Range(_) | Int64(_), Range(_) | Int64(_)) => {
self.to_i64_labels() == other.to_i64_labels()
}
(Datetime(a, _), Datetime(b, _)) => a == b,
(Str(a), Str(b)) => a == b,
_ => false, }
}
pub fn argsort(&self, ascending: bool) -> Vec<usize> {
let mut idx: Vec<usize> = (0..self.len()).collect();
let cmp_dir = |o: std::cmp::Ordering| if ascending { o } else { o.reverse() };
match &self.kind {
IndexKind::Str(v) => idx.sort_by(|&a, &b| cmp_dir(v[a].cmp(&v[b]))),
IndexKind::Datetime(v, _) => idx.sort_by(|&a, &b| {
use std::cmp::Ordering::*;
match (v[a] == i64::MIN, v[b] == i64::MIN) {
(true, true) => Equal,
(true, false) => Greater,
(false, true) => Less,
(false, false) => cmp_dir(v[a].cmp(&v[b])),
}
}),
_ => {
let labels = self.to_i64_labels();
idx.sort_by(|&a, &b| cmp_dir(labels[a].cmp(&labels[b])));
}
}
idx
}
pub fn to_column(&self) -> Column {
match &self.kind {
IndexKind::Range(n) => Column::i64((0..*n as i64).collect()),
IndexKind::Int64(v) => Column::i64(v.clone()),
IndexKind::Datetime(v, _) => Column::datetime(v.clone()),
IndexKind::Str(v) => Column::str(v.clone()),
}
}
pub fn append(&self, other: &Index) -> Result<Index> {
use IndexKind::*;
let kind = match (&self.kind, &other.kind) {
(Range(a), Range(b)) => Range(a + b),
(Datetime(a, ta), Datetime(b, _)) => Datetime([a.as_slice(), b].concat(), *ta),
(Str(a), Str(b)) => Str([a.as_slice(), b].concat()),
(Str(_), _) | (_, Str(_)) => {
return Err(VolasError::Shape(
"cannot append a string index to a non-string index".into(),
))
}
(a, b) => Int64([a.to_i64_labels(), b.to_i64_labels()].concat()),
};
Ok(Index {
name: self.name.clone(),
kind,
})
}
pub fn extend(&mut self, other: &Index) -> Result<()> {
use IndexKind::*;
match (&mut self.kind, &other.kind) {
(Range(a), Range(b)) => *a += b,
(Datetime(a, _), Datetime(b, _)) => a.extend_from_slice(b),
(Int64(a), Int64(b)) => a.extend_from_slice(b),
(Str(a), Str(b)) => a.extend(b.iter().cloned()),
(Str(_), _) | (_, Str(_)) => {
return Err(VolasError::Shape(
"cannot append a string index to a non-string index".into(),
))
}
(slot, b) => {
let mut labels = slot.to_i64_labels();
labels.extend(b.to_i64_labels());
*slot = Int64(labels);
}
}
Ok(())
}
pub fn position_of(&self, label: &Label) -> Option<usize> {
match (&self.kind, label) {
(IndexKind::Range(n), Label::I64(v)) => {
if *v >= 0 && (*v as usize) < *n {
Some(*v as usize)
} else {
None
}
}
(IndexKind::Int64(vs), Label::I64(v)) => vs.iter().position(|x| x == v),
(IndexKind::Datetime(vs, _), Label::I64(v)) => vs.iter().position(|x| x == v),
(IndexKind::Str(vs), Label::Str(s)) => vs.iter().position(|x| x == s),
_ => None,
}
}
pub fn label_slice(&self, lo: Option<&Label>, hi: Option<&Label>) -> (usize, usize) {
match &self.kind {
IndexKind::Str(labels) => {
let start = lo.and_then(Label::as_str).map_or(0, |lo| {
labels
.iter()
.position(|x| x.as_str() >= lo)
.unwrap_or(labels.len())
});
let end = hi.and_then(Label::as_str).map_or(labels.len(), |hi| {
labels
.iter()
.rposition(|x| x.as_str() <= hi)
.map_or(0, |p| p + 1)
});
(start, end.max(start))
}
_ => {
let labels = self.to_i64_labels();
let start = lo.and_then(Label::as_i64).map_or(0, |lo| {
labels.iter().position(|&x| x >= lo).unwrap_or(labels.len())
});
let end = hi.and_then(Label::as_i64).map_or(labels.len(), |hi| {
labels.iter().rposition(|&x| x <= hi).map_or(0, |p| p + 1)
});
(start, end.max(start))
}
}
}
}
fn require_no_missing_labels(col: &Column, kind: &str) -> Result<()> {
if col.null_count() > 0 {
return Err(VolasError::Value(format!(
"cannot use a {kind} column containing volas.NA as an index (a missing \
label has no {kind} representation); drop or fill the NA rows first"
)));
}
Ok(())
}
fn require_unique_labels<T: std::hash::Hash + Eq>(labels: &[T], kind: &str) -> Result<()> {
let mut seen = std::collections::HashSet::with_capacity(labels.len());
for l in labels {
if !seen.insert(l) {
return Err(VolasError::Value(format!(
"cannot use a {kind} column with duplicate labels as an index \
(label access assumes unique labels)"
)));
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn from_datetime_and_int_columns() {
assert_eq!(
Index::from_column(&Column::datetime(vec![5, 6])).unwrap(),
Index::datetime(vec![5, 6], Tz::Naive)
);
assert_eq!(
Index::from_column(&Column::i64(vec![1, 2])).unwrap(),
Index::int64(vec![1, 2])
);
}
#[test]
fn from_unsupported_column_errors() {
assert!(Index::from_column(&Column::f64(vec![1.0])).is_err());
assert!(Index::from_column(&Column::bool(vec![true])).is_err());
}
#[test]
fn is_empty_labels_and_position_of() {
assert!(Index::range(0).is_empty());
assert!(!Index::range(3).is_empty());
assert_eq!(Index::range(3).to_i64_labels(), vec![0, 1, 2]);
assert_eq!(Index::int64(vec![5, 6]).to_i64_labels(), vec![5, 6]);
assert_eq!(
Index::datetime(vec![10, 20], Tz::Utc).to_i64_labels(),
vec![10, 20]
);
let i64 = Label::I64;
assert_eq!(Index::range(5).position_of(&i64(3)), Some(3));
assert_eq!(Index::range(5).position_of(&i64(9)), None);
assert_eq!(Index::range(5).position_of(&i64(-1)), None);
assert_eq!(Index::int64(vec![10, 20, 30]).position_of(&i64(20)), Some(1));
assert_eq!(Index::int64(vec![10, 20]).position_of(&i64(99)), None);
assert_eq!(
Index::datetime(vec![100, 200], Tz::Utc).position_of(&i64(200)),
Some(1)
);
assert_eq!(
Index::int64(vec![10, 20, 30]).take(&[2, 0]),
Index::int64(vec![30, 10])
);
}
fn str_index(labels: &[&str]) -> Index {
Index::str(labels.iter().map(|s| s.to_string()).collect())
}
#[test]
fn string_index_construction_and_ops() {
let ix = Index::from_column(&Column::str(vec!["a".into(), "b".into()])).unwrap();
assert_eq!(ix, str_index(&["a", "b"]));
let ix = str_index(&["a", "b", "c", "d"]);
assert_eq!(ix.len(), 4);
assert!(!ix.is_empty());
assert_eq!(ix.label_at(2), Label::Str("c".into()));
assert_eq!(ix.slice(1, 3), str_index(&["b", "c"]));
assert_eq!(ix.take(&[3, 0]), str_index(&["d", "a"]));
}
#[test]
fn string_index_lookup_and_slice() {
let ix = str_index(&["aa", "bb", "cc", "dd"]);
assert_eq!(ix.position_of(&Label::Str("cc".into())), Some(2));
assert_eq!(ix.position_of(&Label::Str("zz".into())), None);
assert_eq!(ix.position_of(&Label::I64(1)), None);
let lo = Label::Str("bb".into());
let hi = Label::Str("cc".into());
assert_eq!(ix.label_slice(Some(&lo), Some(&hi)), (1, 3));
assert_eq!(ix.label_slice(Some(&lo), None), (1, 4));
}
#[test]
fn string_index_append_rules() {
let a = str_index(&["x", "y"]);
let b = str_index(&["z"]);
assert_eq!(a.append(&b).unwrap(), str_index(&["x", "y", "z"]));
assert!(a.append(&Index::range(2)).is_err());
assert!(Index::range(2).append(&a).is_err());
}
#[test]
fn extend_grows_in_place_per_kind() {
let mut r = Index::range(3);
r.extend(&Index::range(2)).unwrap();
assert_eq!(r, Index::range(5));
let mut d = Index::datetime(vec![1, 2], Tz::Utc);
d.extend(&Index::datetime(vec![3], Tz::Utc)).unwrap();
assert_eq!(d, Index::datetime(vec![1, 2, 3], Tz::Utc));
let mut s = str_index(&["a", "b"]);
s.extend(&str_index(&["c"])).unwrap();
assert_eq!(s, str_index(&["a", "b", "c"]));
let mut m = Index::range(2);
m.extend(&Index::int64(vec![5, 6])).unwrap();
assert_eq!(m, Index::int64(vec![0, 1, 5, 6]));
assert!(str_index(&["x"]).extend(&Index::range(1)).is_err());
assert!(Index::range(1).extend(&str_index(&["x"])).is_err());
}
#[test]
fn name_set_and_propagates_through_ops() {
let ix = Index::datetime(vec![1, 2, 3], Tz::Utc).with_name(Some("date".into()));
assert_eq!(ix.name(), Some("date"));
assert_eq!(Index::range(3).name(), None);
assert_eq!(ix.slice(0, 2).name(), Some("date"));
assert_eq!(ix.take(&[2, 0]).name(), Some("date"));
assert_eq!(ix.clone().with_tz(Tz::Offset(28800)).name(), Some("date"));
assert_eq!(
ix.append(&Index::datetime(vec![4], Tz::Utc)).unwrap().name(),
Some("date")
);
let mut g = ix.clone();
g.extend(&Index::datetime(vec![4], Tz::Utc)).unwrap();
assert_eq!(g.name(), Some("date"));
assert_eq!(ix.with_name(None).name(), None);
}
#[test]
fn label_accessors_and_numeric_label_at() {
assert_eq!(Label::I64(5).as_i64(), Some(5));
assert_eq!(Label::I64(5).as_str(), None);
assert_eq!(Label::Str("x".into()).as_str(), Some("x"));
assert_eq!(Label::Str("x".into()).as_i64(), None);
assert_eq!(Index::range(3).label_at(2), Label::I64(2));
assert_eq!(Index::int64(vec![10, 20]).label_at(1), Label::I64(20));
assert_eq!(
Index::datetime(vec![100, 200], Tz::Utc).label_at(0),
Label::I64(100)
);
}
#[test]
fn index_kind_branch_coverage() {
assert_eq!(Index::range(3).tz(), Tz::Naive);
assert!(matches!(
Index::range(3).with_tz(Tz::Utc).kind,
IndexKind::Range(3)
));
assert_eq!(
Index::int64(vec![1, 2, 3]).slice(0, 2),
Index::int64(vec![1, 2])
);
assert!(matches!(
Index::datetime(vec![1, 2], Tz::Utc).slice(0, 1).kind,
IndexKind::Datetime(_, _)
));
assert_eq!(
Index::str(vec!["a".into(), "b".into()]).slice(1, 2),
Index::str(vec!["b".into()])
);
assert_eq!(
Index::str(vec!["b".into(), "a".into()]).argsort(true),
vec![1, 0]
);
assert_eq!(Index::range(2).to_column().len(), 2);
assert_eq!(Index::datetime(vec![5], Tz::Utc).to_column().len(), 1);
assert_eq!(Index::str(vec!["x".into()]).to_column().len(), 1);
assert!(matches!(
Index::datetime(vec![1], Tz::Utc)
.append(&Index::datetime(vec![2], Tz::Utc))
.unwrap()
.kind,
IndexKind::Datetime(_, _)
));
assert!(matches!(
Index::str(vec!["a".into()])
.append(&Index::str(vec!["b".into()]))
.unwrap()
.kind,
IndexKind::Str(_)
));
assert!(matches!(
Index::range(2).append(&Index::range(3)).unwrap().kind,
IndexKind::Range(5)
));
assert!(matches!(
Index::range(2).append(&Index::int64(vec![5])).unwrap().kind,
IndexKind::Int64(_)
));
assert!(Index::str(vec!["a".into()])
.append(&Index::range(1))
.is_err());
}
#[test]
fn label_eq_value_semantics() {
assert!(Index::range(3).label_eq(&Index::int64(vec![0, 1, 2])));
assert!(!Index::range(3).label_eq(&Index::int64(vec![0, 1, 9])));
assert!(Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 2], Tz::Utc)));
assert!(!Index::datetime(vec![1, 2], Tz::Utc).label_eq(&Index::datetime(vec![1, 9], Tz::Utc)));
assert!(Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["a".into()])));
assert!(!Index::str(vec!["a".into()]).label_eq(&Index::str(vec!["b".into()])));
assert!(!Index::range(2).label_eq(&Index::datetime(vec![0, 1], Tz::Utc)));
assert!(!Index::str(vec!["a".into()]).label_eq(&Index::int64(vec![0])));
}
}