use super::selectors::parse_variable_selector;
use super::*;
mod grpstats;
pub(in crate::builtins::table) use grpstats::grpstats_impl;
pub fn sortrows_table(value: Value, rest: &[Value]) -> BuiltinResult<(Value, Tensor)> {
let object = into_table_object(value, "sortrows")?;
let names = table_variable_names_from_object(&object)?;
let sort_spec = SortSpec::parse(rest, &names)?;
let height = table_height(&object)?;
let variables = table_variables(&object)?;
let mut indices: Vec<usize> = (0..height).collect();
indices.sort_by(|&a, &b| {
for key in &sort_spec.keys {
let Some(value) = variables.fields.get(&key.name) else {
continue;
};
let ord = compare_table_cells(value, a, b).unwrap_or(Ordering::Equal);
let ord = if key.descending { ord.reverse() } else { ord };
if ord != Ordering::Equal {
return ord;
}
}
a.cmp(&b)
});
let mut sorted_columns = Vec::with_capacity(names.len());
for name in &names {
let value = variables
.fields
.get(name)
.ok_or_else(|| invalid_variable(format!("table: missing variable '{name}'")))?;
sorted_columns.push(select_rows(value, &indices)?);
}
let row_names = selected_row_names(&object, &indices)?;
let sorted = table_from_columns_with_properties(names, sorted_columns, row_names)?;
let indices_tensor = Tensor::new(
indices.iter().map(|idx| *idx as f64 + 1.0).collect(),
vec![indices.len(), 1],
)
.map_err(invalid_variable)?;
Ok((sorted, indices_tensor))
}
pub(in crate::builtins::table) struct SortSpec {
keys: Vec<SortKey>,
}
pub(in crate::builtins::table) struct SortKey {
name: String,
descending: bool,
}
impl SortSpec {
fn parse(rest: &[Value], names: &[String]) -> BuiltinResult<Self> {
let mut keys = if rest.is_empty() {
names
.iter()
.map(|name| SortKey {
name: name.clone(),
descending: false,
})
.collect::<Vec<_>>()
} else {
parse_variable_selector(rest.first(), names)?
.into_iter()
.map(|name| SortKey {
name,
descending: false,
})
.collect()
};
if let Some(direction) = rest.get(1) {
let directions = string_list(direction)?;
if directions.len() == 1 {
let descending = directions[0].eq_ignore_ascii_case("descend")
|| directions[0].eq_ignore_ascii_case("desc");
for key in &mut keys {
key.descending = descending;
}
} else {
for (key, direction) in keys.iter_mut().zip(directions.iter()) {
key.descending = direction.eq_ignore_ascii_case("descend")
|| direction.eq_ignore_ascii_case("desc");
}
}
}
Ok(Self { keys })
}
}
pub(in crate::builtins::table) fn compare_table_cells(
value: &Value,
a: usize,
b: usize,
) -> BuiltinResult<Ordering> {
match value {
Value::Tensor(tensor) => Ok(tensor
.get2(a, 0)
.map_err(invalid_index)?
.partial_cmp(&tensor.get2(b, 0).map_err(invalid_index)?)
.unwrap_or(Ordering::Greater)),
Value::StringArray(array) => {
let av = array.data.get(a).cloned().unwrap_or_default();
let bv = array.data.get(b).cloned().unwrap_or_default();
Ok(av.cmp(&bv))
}
Value::LogicalArray(array) => {
let av = *array.data.get(a).unwrap_or(&0);
let bv = *array.data.get(b).unwrap_or(&0);
Ok(av.cmp(&bv))
}
Value::Object(obj) if obj.is_class("datetime") => {
let tensor = crate::builtins::datetime::serials_from_datetime_value(value)?;
Ok(tensor
.data
.get(a)
.copied()
.unwrap_or(f64::NAN)
.partial_cmp(&tensor.data.get(b).copied().unwrap_or(f64::NAN))
.unwrap_or(Ordering::Greater))
}
other => Ok(cell_key_string(other, a).cmp(&cell_key_string(other, b))),
}
}
#[derive(Clone, Debug)]
pub(in crate::builtins::table) enum GroupAtom {
Number(f64),
Text(String),
Logical(bool),
Missing,
}
impl GroupAtom {
fn rank(&self) -> u8 {
match self {
Self::Missing => 0,
Self::Logical(_) => 1,
Self::Number(_) => 2,
Self::Text(_) => 3,
}
}
}
impl PartialEq for GroupAtom {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for GroupAtom {}
impl PartialOrd for GroupAtom {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for GroupAtom {
fn cmp(&self, other: &Self) -> Ordering {
let rank = self.rank().cmp(&other.rank());
if rank != Ordering::Equal {
return rank;
}
match (self, other) {
(Self::Missing, Self::Missing) => Ordering::Equal,
(Self::Logical(a), Self::Logical(b)) => a.cmp(b),
(Self::Number(a), Self::Number(b)) => a.total_cmp(b),
(Self::Text(a), Self::Text(b)) => a.cmp(b),
_ => Ordering::Equal,
}
}
}
pub(in crate::builtins::table) fn cell_group_atom(value: &Value, row: usize) -> GroupAtom {
match value {
Value::Tensor(tensor) => tensor
.get2(row, 0)
.map(GroupAtom::Number)
.unwrap_or(GroupAtom::Missing),
Value::StringArray(array) => array
.data
.get(row)
.cloned()
.map(GroupAtom::Text)
.unwrap_or(GroupAtom::Missing),
Value::LogicalArray(array) => array
.data
.get(row)
.map(|value| GroupAtom::Logical(*value != 0))
.unwrap_or(GroupAtom::Missing),
Value::Object(obj) if obj.is_class("datetime") => {
crate::builtins::datetime::serials_from_datetime_value(value)
.ok()
.and_then(|tensor| tensor.data.get(row).copied())
.map(GroupAtom::Number)
.unwrap_or(GroupAtom::Missing)
}
other => GroupAtom::Text(cell_key_string(other, row)),
}
}
pub(in crate::builtins::table) fn pivot_impl(
table: Value,
rowvars: Value,
colvars: Value,
datavar: Value,
method: &str,
) -> BuiltinResult<Value> {
let object = into_table_object(table, "pivot")?;
let names = table_variable_names_from_object(&object)?;
let row_names = parse_variable_selector_for_object(Some(&rowvars), &object, &names)?;
let col_names = parse_variable_selector_for_object(Some(&colvars), &object, &names)?;
let data_names = parse_variable_selector_for_object(Some(&datavar), &object, &names)?;
if row_names.is_empty() || col_names.is_empty() || data_names.is_empty() {
return Err(invalid_argument(
"pivot: rowvars, colvars, and datavar must select at least one variable",
));
}
if data_names.len() != 1 {
return Err(invalid_argument(
"pivot: exactly one data variable is currently supported",
));
}
let data_name = &data_names[0];
let variables = table_variables(&object)?;
let data_value = variables
.fields
.get(data_name)
.ok_or_else(|| invalid_variable(format!("pivot: missing data variable '{data_name}'")))?;
if !matches!(data_value, Value::Tensor(tensor) if tensor.cols() == 1) {
return Err(invalid_variable(
"pivot: data variable must be a numeric column vector",
));
}
let height = table_height(&object)?;
let mut row_order = Vec::<Vec<GroupAtom>>::new();
let mut row_first_index = BTreeMap::<Vec<GroupAtom>, usize>::new();
let mut col_order = Vec::<Vec<GroupAtom>>::new();
let mut col_seen = BTreeMap::<Vec<GroupAtom>, ()>::new();
let mut buckets = BTreeMap::<(Vec<GroupAtom>, Vec<GroupAtom>), Vec<usize>>::new();
for row in 0..height {
let row_key = group_key_for_row(&variables, &row_names, row);
let col_key = group_key_for_row(&variables, &col_names, row);
if !row_first_index.contains_key(&row_key) {
row_first_index.insert(row_key.clone(), row);
row_order.push(row_key.clone());
}
if !col_seen.contains_key(&col_key) {
col_seen.insert(col_key.clone(), ());
col_order.push(col_key.clone());
}
buckets.entry((row_key, col_key)).or_default().push(row);
}
let mut out_names = row_names.clone();
let mut out_columns = Vec::with_capacity(row_names.len() + col_order.len());
for name in &row_names {
let value = variables
.fields
.get(name)
.ok_or_else(|| invalid_variable(format!("pivot: missing row variable '{name}'")))?;
let rows = row_order
.iter()
.filter_map(|key| row_first_index.get(key).copied())
.collect::<Vec<_>>();
out_columns.push(select_rows(value, &rows)?);
}
for col_key in &col_order {
let mut values = Vec::with_capacity(row_order.len());
for row_key in &row_order {
let summary_rows = buckets
.get(&(row_key.clone(), col_key.clone()))
.cloned()
.unwrap_or_default();
if summary_rows.is_empty() {
values.push(f64::NAN);
} else {
values.push(
summarize_groups(data_value, std::iter::once(&summary_rows), method)?
.into_iter()
.next()
.unwrap_or(f64::NAN),
);
}
}
out_names.push(format!(
"{}_{}",
make_valid_variable_name(&group_key_label(col_key), out_names.len() + 1),
data_name
));
out_columns.push(Value::Tensor(
Tensor::new(values, vec![row_order.len(), 1]).map_err(invalid_variable)?,
));
}
let out_names = make_unique_variable_names(out_names);
table_from_columns(out_names, out_columns)
}
pub(in crate::builtins::table) fn group_key_for_row(
variables: &StructValue,
names: &[String],
row: usize,
) -> Vec<GroupAtom> {
names
.iter()
.map(|name| {
variables
.fields
.get(name)
.map(|value| cell_group_atom(value, row))
.unwrap_or(GroupAtom::Missing)
})
.collect()
}
pub(in crate::builtins::table) fn group_key_label(key: &[GroupAtom]) -> String {
if key.is_empty() {
return "missing".to_string();
}
key.iter()
.map(group_atom_label)
.collect::<Vec<_>>()
.join("_")
}
pub(in crate::builtins::table) fn group_atom_label(atom: &GroupAtom) -> String {
match atom {
GroupAtom::Number(value) => format_key_number(*value),
GroupAtom::Text(text) => text.clone(),
GroupAtom::Logical(flag) => flag.to_string(),
GroupAtom::Missing => "missing".to_string(),
}
}
pub(in crate::builtins::table) fn groupsummary_impl(
table: Value,
groupvars: Value,
method: Value,
rest: Vec<Value>,
) -> BuiltinResult<Value> {
let object = into_table_object(table, "groupsummary")?;
let names = table_variable_names_from_object(&object)?;
let group_names = parse_variable_selector_for_object(Some(&groupvars), &object, &names)?;
let methods = string_list(&method)?;
if methods.is_empty() {
return Err(invalid_argument(
"groupsummary: method list must not be empty",
));
}
let data_names = if let Some(value) = rest.first() {
parse_variable_selector_for_object(Some(value), &object, &names)?
} else {
names
.iter()
.filter(|name| !group_names.contains(name))
.filter(|name| {
table_variables(&object)
.ok()
.and_then(|vars| vars.fields.get(*name).cloned())
.map(|value| matches!(value, Value::Tensor(_)))
.unwrap_or(false)
})
.cloned()
.collect()
};
let variables = table_variables(&object)?;
let height = table_height(&object)?;
let mut groups: BTreeMap<Vec<GroupAtom>, Vec<usize>> = BTreeMap::new();
for row in 0..height {
let key = group_names
.iter()
.map(|name| {
variables
.fields
.get(name)
.map(|value| cell_group_atom(value, row))
.unwrap_or(GroupAtom::Missing)
})
.collect::<Vec<_>>();
groups.entry(key).or_default().push(row);
}
let group_rows = groups
.values()
.filter_map(|rows| rows.first().copied())
.collect::<Vec<_>>();
let mut out_names = Vec::new();
let mut out_columns = Vec::new();
for name in &group_names {
let value = variables.fields.get(name).ok_or_else(|| {
invalid_variable(format!("groupsummary: missing group variable '{name}'"))
})?;
out_names.push(name.clone());
out_columns.push(select_rows(value, &group_rows)?);
}
out_names.push("GroupCount".to_string());
out_columns.push(Value::Tensor(
Tensor::new(
groups.values().map(|rows| rows.len() as f64).collect(),
vec![groups.len(), 1],
)
.map_err(invalid_variable)?,
));
for method in &methods {
for name in &data_names {
let value = variables.fields.get(name).ok_or_else(|| {
invalid_variable(format!("groupsummary: missing data variable '{name}'"))
})?;
let values = summarize_groups(value, groups.values(), method)?;
out_names.push(format!("{}_{}", method.to_ascii_lowercase(), name));
out_columns.push(Value::Tensor(
Tensor::new(values, vec![groups.len(), 1]).map_err(invalid_variable)?,
));
}
}
table_from_columns(out_names, out_columns)
}
pub(in crate::builtins::table) fn summarize_groups<'a>(
value: &Value,
groups: impl Iterator<Item = &'a Vec<usize>>,
method: &str,
) -> BuiltinResult<Vec<f64>> {
let tensor = match value {
Value::Tensor(tensor) if tensor.cols() == 1 => tensor,
_ => {
return Err(invalid_variable(
"groupsummary: summary data variables must be numeric column vectors",
))
}
};
groups
.map(|rows| {
let mut values = rows
.iter()
.map(|row| tensor.get2(*row, 0).map_err(invalid_index))
.collect::<BuiltinResult<Vec<_>>>()?;
values.retain(|value| !value.is_nan());
let result = match method.to_ascii_lowercase().as_str() {
"mean" => {
if values.is_empty() {
f64::NAN
} else {
values.iter().sum::<f64>() / values.len() as f64
}
}
"sum" => values.iter().sum(),
"min" => values.into_iter().fold(f64::INFINITY, f64::min),
"max" => values.into_iter().fold(f64::NEG_INFINITY, f64::max),
"median" => {
if values.is_empty() {
f64::NAN
} else {
values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal));
let mid = values.len() / 2;
if values.len() % 2 == 0 {
(values[mid - 1] + values[mid]) / 2.0
} else {
values[mid]
}
}
}
"count" | "numel" => values.len() as f64,
other => {
return Err(invalid_argument(format!(
"groupsummary: unsupported method '{other}'"
)))
}
};
Ok(result)
})
.collect()
}
pub(in crate::builtins::table) fn cell_key_string(value: &Value, row: usize) -> String {
match value {
Value::Tensor(tensor) => tensor
.get2(row, 0)
.map(format_key_number)
.unwrap_or_default(),
Value::StringArray(array) => array.data.get(row).cloned().unwrap_or_default(),
Value::LogicalArray(array) => array
.data
.get(row)
.map(|value| value.to_string())
.unwrap_or_default(),
Value::Object(obj) if obj.is_class("datetime") => {
crate::builtins::datetime::serials_from_datetime_value(value)
.ok()
.and_then(|tensor| tensor.data.get(row).copied())
.map(format_key_number)
.unwrap_or_default()
}
Value::Object(obj) if obj.is_class("duration") => {
crate::builtins::duration::duration_tensor_from_duration_value(value)
.ok()
.and_then(|tensor| tensor.data.get(row).copied())
.map(format_key_number)
.unwrap_or_default()
}
Value::Object(obj) if obj.is_class(CATEGORICAL_CLASS) => {
categorical_label_at(obj, row).unwrap_or_default()
}
Value::Cell(cell) => cell
.get(row, 0)
.map(|item| cell_to_text(&item))
.unwrap_or_default(),
other => format!("{other}"),
}
}