use std::cmp::Ordering;
use std::collections::BTreeMap;
use runmat_builtins::{
BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
CellArray, LogicalArray, ObjectInstance, SparseTensor, StringArray, Tensor, Value,
};
use runmat_macros::runtime_builtin;
use crate::builtins::table::{
categorical_label_at, is_tabular_object, select_rows, table_from_columns, table_height,
table_variable_names_from_object, table_variables, value_row_count,
};
use crate::{
build_runtime_error, call_feval_async_with_outputs, gather_if_needed_async, BuiltinResult,
RuntimeError,
};
const MAX_MATERIALIZED_ELEMENTS: usize = 50_000_000;
const OUTPUT_ANY: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "varargout",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Variadic,
default: None,
description: "Builtin outputs.",
}];
const INPUT_VARIADIC: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
name: "args",
ty: BuiltinParamType::Any,
arity: BuiltinParamArity::Variadic,
default: None,
description: "MATLAB-compatible arguments.",
}];
const SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
label: "varargout = groupingBuiltin(args...)",
inputs: &INPUT_VARIADIC,
outputs: &OUTPUT_ANY,
}];
const ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.GROUPING.INVALID_INPUT",
identifier: Some("RunMat:grouping:InvalidInput"),
when: "Inputs are malformed, have incompatible lengths, or request unsupported grouped output.",
message: "grouping builtin: invalid input",
};
const ERROR_CALLBACK: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.GROUPING.CALLBACK",
identifier: Some("RunMat:grouping:CallbackFailed"),
when: "A grouped callback fails or returns incompatible outputs.",
message: "grouping builtin: callback failed",
};
const ERROR_TOO_LARGE: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
code: "RM.GROUPING.TOO_LARGE",
identifier: Some("RunMat:grouping:TooLarge"),
when: "The requested dense or combinatorial output exceeds RunMat's materialization limit.",
message: "grouping builtin: output is too large",
};
const ERRORS: [BuiltinErrorDescriptor; 3] = [ERROR_INVALID_INPUT, ERROR_CALLBACK, ERROR_TOO_LARGE];
pub const GROUPING_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
signatures: &SIGNATURES,
output_mode: BuiltinOutputMode::ByRequestedOutputCount,
completion_policy: BuiltinCompletionPolicy::Public,
errors: &ERRORS,
};
#[derive(Clone, Debug)]
enum Atom {
Missing,
Logical(bool),
Number(f64),
Text(String),
}
impl Atom {
fn rank(&self) -> u8 {
match self {
Self::Missing => 0,
Self::Logical(_) => 1,
Self::Number(_) => 2,
Self::Text(_) => 3,
}
}
fn label(&self) -> String {
match self {
Self::Missing => "<missing>".to_string(),
Self::Logical(flag) => {
if *flag {
"true".to_string()
} else {
"false".to_string()
}
}
Self::Number(value) => format_key_number(*value),
Self::Text(text) => text.clone(),
}
}
}
impl PartialEq for Atom {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for Atom {}
impl PartialOrd for Atom {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Atom {
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,
}
}
}
#[derive(Clone)]
struct GroupColumn {
name: String,
value: Value,
rows: usize,
}
struct Grouping {
ids: Vec<f64>,
keys: Vec<Vec<Atom>>,
first_rows: Vec<usize>,
row_groups: Vec<Vec<usize>>,
}
#[derive(Clone, Copy, Debug)]
struct GroupOptions {
include_missing: bool,
}
impl GroupOptions {
fn parse(args: &[Value], context: &str) -> BuiltinResult<Self> {
let mut include_missing = false;
let mut idx = 0usize;
while idx < args.len() {
if idx + 1 >= args.len() {
return Err(grouping_error(format!(
"{context}: name-value options must be provided in pairs"
)));
}
let name = scalar_text(&args[idx], context)?;
if name.eq_ignore_ascii_case("IncludeMissingGroups") {
include_missing = bool_scalar(&args[idx + 1], "IncludeMissingGroups")?;
} else if name.eq_ignore_ascii_case("IncludeEmptyGroups") {
let include_empty = bool_scalar(&args[idx + 1], "IncludeEmptyGroups")?;
if include_empty {
return Err(grouping_error(format!(
"{context}: IncludeEmptyGroups=true is not supported until categorical level expansion is implemented"
)));
}
} else {
return Err(grouping_error(format!(
"{context}: unsupported option '{name}'"
)));
}
idx += 2;
}
Ok(Self { include_missing })
}
}
fn grouping_error(message: impl Into<String>) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin("array_grouping");
if let Some(identifier) = ERROR_INVALID_INPUT.identifier {
builder = builder.with_identifier(identifier);
}
builder.build()
}
fn callback_error(message: impl Into<String>, source: Option<RuntimeError>) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin("array_grouping");
if let Some(identifier) = ERROR_CALLBACK.identifier {
builder = builder.with_identifier(identifier);
}
if let Some(source) = source {
builder = builder.with_source(source);
}
builder.build()
}
fn too_large_error(message: impl Into<String>) -> RuntimeError {
let mut builder = build_runtime_error(message).with_builtin("array_grouping");
if let Some(identifier) = ERROR_TOO_LARGE.identifier {
builder = builder.with_identifier(identifier);
}
builder.build()
}
#[runtime_builtin(
name = "findgroups",
category = "array/grouping",
summary = "Find groups and return group numbers.",
keywords = "findgroups,groups,grouping,table,categorical",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn findgroups_builtin(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let mut args = Vec::with_capacity(rest.len() + 1);
args.push(gather_if_needed_async(&first).await?);
for value in rest {
args.push(gather_if_needed_async(&value).await?);
}
let (columns, table_mode) = findgroups_columns(args)?;
let grouping = build_grouping(&columns)?;
let outputs = findgroups_outputs(&columns, &grouping, table_mode)?;
multi_output(outputs)
}
#[runtime_builtin(
name = "grp2idx",
category = "array/grouping",
summary = "Create an index vector from a grouping variable.",
keywords = "grp2idx,groups,index,categorical,statistics",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn grp2idx_builtin(value: Value) -> BuiltinResult<Value> {
let value = gather_if_needed_async(&value).await?;
let columns = columns_from_group_value("G", value, true)?;
if columns.len() != 1 {
return Err(grouping_error(
"grp2idx: expected one grouping vector, not a matrix of grouping columns",
));
}
let grouping = build_grouping(&columns)?;
let g = Tensor::new(grouping.ids.clone(), vec![grouping.ids.len(), 1])
.map(Value::Tensor)
.map_err(grouping_error)?;
let names = grouping
.keys
.iter()
.map(|key| key.first().map(Atom::label).unwrap_or_default())
.collect::<Vec<_>>();
let gn = Value::StringArray(
StringArray::new(names.clone(), vec![names.len(), 1]).map_err(grouping_error)?,
);
let gl = Value::StringArray(
StringArray::new(names, vec![grouping.keys.len(), 1]).map_err(grouping_error)?,
);
multi_output(vec![g, gn, gl])
}
#[runtime_builtin(
name = "groupcounts",
category = "array/grouping",
summary = "Count the number of elements in each group.",
keywords = "groupcounts,groups,count,table,categorical",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn groupcounts_builtin(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let first = gather_if_needed_async(&first).await?;
let rest = gather_values(rest).await?;
if let Value::Object(object) = first.clone() {
if is_tabular_object(&object) {
let (selector_args, option_args) = split_option_tail(rest)?;
return groupcounts_table(object, selector_args, option_args);
}
}
let (data_args, option_args) = split_option_tail(rest)?;
let options = GroupOptions::parse(&option_args, "groupcounts")?;
let columns = if matches!(first, Value::Cell(_)) {
columns_from_group_value("A", first, true)?
} else {
let mut args = vec![first];
args.extend(data_args);
columns_from_group_args(args)?
};
let grouping = build_grouping_with_options(&columns, options)?;
let counts = grouping
.row_groups
.iter()
.map(|rows| rows.len() as f64)
.collect::<Vec<_>>();
let b =
Value::Tensor(Tensor::new(counts, vec![grouping.keys.len(), 1]).map_err(grouping_error)?);
let bg = group_label_outputs(&columns, &grouping)?;
let bp = Value::Tensor(
Tensor::new(
grouping
.row_groups
.iter()
.map(|rows| {
if grouping.ids.is_empty() {
0.0
} else {
rows.len() as f64 * 100.0 / grouping.ids.len() as f64
}
})
.collect(),
vec![grouping.keys.len(), 1],
)
.map_err(grouping_error)?,
);
let mut outputs = vec![b];
outputs.extend(bg);
outputs.push(bp);
multi_output(outputs)
}
#[runtime_builtin(
name = "splitapply",
category = "array/grouping",
summary = "Split data into groups and apply a function.",
keywords = "splitapply,groups,apply,function,table",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn splitapply_builtin(
func: Value,
first_data: Value,
rest: Vec<Value>,
) -> BuiltinResult<Value> {
let func = gather_if_needed_async(&func).await?;
let first_data = gather_if_needed_async(&first_data).await?;
let rest = gather_values(rest).await?;
if rest.is_empty() {
return Err(grouping_error(
"splitapply: expected data arguments followed by group numbers",
));
}
let (group_value, data_tail) = rest.split_last().expect("checked non-empty");
let mut data_values = Vec::with_capacity(data_tail.len() + 1);
data_values.push(first_data);
data_values.extend_from_slice(data_tail);
let group_columns = columns_from_group_value("G", group_value.clone(), true)?;
if group_columns.len() != 1 {
return Err(grouping_error("splitapply: G must be a grouping vector"));
}
let grouping = build_grouping(&group_columns)?;
let expected_rows = grouping.ids.len();
for value in &data_values {
let rows = value_row_count(value)?;
if rows != expected_rows {
return Err(grouping_error(format!(
"splitapply: data arguments must have {expected_rows} rows to match G, got {rows}"
)));
}
}
let requested_outputs = crate::output_count::current_output_count()
.unwrap_or(1)
.max(1);
let mut collectors = (0..requested_outputs)
.map(|_| Vec::<Value>::new())
.collect::<Vec<_>>();
for rows in &grouping.row_groups {
let mut callback_args = Vec::with_capacity(data_values.len());
for value in &data_values {
callback_args.push(select_rows(value, rows)?);
}
let result = call_feval_async_with_outputs(func.clone(), &callback_args, requested_outputs)
.await
.map_err(|err| callback_error("splitapply: callback failed", Some(err)))?;
let outputs = normalize_outputs(result, requested_outputs, "splitapply")?;
for (collector, output) in collectors.iter_mut().zip(outputs) {
collector.push(gather_if_needed_async(&output).await?);
}
}
let outputs = collectors
.into_iter()
.map(|values| collect_group_results(values, grouping.keys.len(), "splitapply"))
.collect::<BuiltinResult<Vec<_>>>()?;
multi_output(outputs)
}
#[runtime_builtin(
name = "accumarray",
category = "array/grouping",
summary = "Accumulate values into an array by subscript groups.",
keywords = "accumarray,accumulate,groups,sum,sparse",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn accumarray_builtin(
subs: Value,
data: Value,
rest: Vec<Value>,
) -> BuiltinResult<Value> {
let subs = gather_if_needed_async(&subs).await?;
let data = gather_if_needed_async(&data).await?;
let rest = gather_values(rest).await?;
accumarray_impl(subs, data, rest).await
}
#[runtime_builtin(
name = "discretize",
category = "array/grouping",
summary = "Group numeric data into bins.",
keywords = "discretize,bins,edges,categorical,grouping",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn discretize_builtin(
x: Value,
edges_or_n: Value,
rest: Vec<Value>,
) -> BuiltinResult<Value> {
let x = gather_if_needed_async(&x).await?;
let edges_or_n = gather_if_needed_async(&edges_or_n).await?;
let rest = gather_values(rest).await?;
discretize_impl(x, edges_or_n, rest)
}
#[runtime_builtin(
name = "combinations",
category = "array/grouping",
summary = "Generate all element combinations of arrays.",
keywords = "combinations,cartesian,table,combinatorics",
accel = "cpu",
descriptor(crate::builtins::array::grouping::GROUPING_DESCRIPTOR),
builtin_path = "crate::builtins::array::grouping"
)]
pub(crate) async fn combinations_builtin(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let first = gather_if_needed_async(&first).await?;
let rest = gather_values(rest).await?;
combinations_impl(first, rest)
}
fn findgroups_columns(args: Vec<Value>) -> BuiltinResult<(Vec<GroupColumn>, bool)> {
if args.is_empty() {
return Err(grouping_error("findgroups: expected at least one input"));
}
if let Value::Object(object) = args[0].clone() {
if is_tabular_object(&object) {
let names = table_variable_names_from_object(&object)?;
let selected = if let Some(selector) = args.get(1) {
parse_name_selector(selector, &names, "findgroups")?
} else {
names
};
let variables = table_variables(&object)?;
let mut columns = Vec::with_capacity(selected.len());
for name in selected {
let value = variables.fields.get(&name).cloned().ok_or_else(|| {
grouping_error(format!("findgroups: missing table variable '{name}'"))
})?;
columns.push(GroupColumn {
rows: value_row_count(&value)?,
name,
value,
});
}
return Ok((columns, true));
}
}
Ok((columns_from_group_args(args)?, false))
}
fn columns_from_group_args(args: Vec<Value>) -> BuiltinResult<Vec<GroupColumn>> {
let mut columns = Vec::new();
for (idx, value) in args.into_iter().enumerate() {
columns.extend(columns_from_group_value(
&format!("Var{}", idx + 1),
value,
true,
)?);
}
Ok(columns)
}
fn columns_from_group_value(
base_name: &str,
value: Value,
split_matrix: bool,
) -> BuiltinResult<Vec<GroupColumn>> {
match value {
Value::Tensor(tensor) => tensor_columns(base_name, tensor, split_matrix),
Value::LogicalArray(array) => logical_columns(base_name, array, split_matrix),
Value::StringArray(array) => string_columns(base_name, array, split_matrix),
Value::Cell(cell) if cell_is_group_vector_list(&cell) => {
let mut columns = Vec::with_capacity(cell.data.len());
for (idx, value) in cell.data.into_iter().enumerate() {
columns.extend(columns_from_group_value(
&format!("{base_name}{}", idx + 1),
value,
false,
)?);
}
Ok(columns)
}
Value::Cell(cell) => Ok(vec![GroupColumn {
rows: cell.rows.max(cell.cols).max(cell.data.len()),
name: base_name.to_string(),
value: Value::Cell(cell),
}]),
Value::Object(object) if object.is_class("categorical") => {
let rows = value_row_count(&Value::Object(object.clone()))?;
Ok(vec![GroupColumn {
rows,
name: base_name.to_string(),
value: Value::Object(object),
}])
}
other => {
let rows = value_row_count(&other).unwrap_or(1);
Ok(vec![GroupColumn {
rows,
name: base_name.to_string(),
value: other,
}])
}
}
}
fn tensor_columns(
base_name: &str,
tensor: Tensor,
split_matrix: bool,
) -> BuiltinResult<Vec<GroupColumn>> {
if !split_matrix || tensor.cols() <= 1 || tensor.rows() == 1 {
let rows = tensor.data.len();
return Ok(vec![GroupColumn {
name: base_name.to_string(),
rows,
value: Value::Tensor(Tensor::new(tensor.data, vec![rows, 1]).map_err(grouping_error)?),
}]);
}
let mut out = Vec::with_capacity(tensor.cols());
for col in 0..tensor.cols() {
let mut data = Vec::with_capacity(tensor.rows());
for row in 0..tensor.rows() {
data.push(tensor.get2(row, col).map_err(grouping_error)?);
}
out.push(GroupColumn {
name: format!("{base_name}{col_plus}", col_plus = col + 1),
rows: tensor.rows(),
value: Value::Tensor(
Tensor::new(data, vec![tensor.rows(), 1]).map_err(grouping_error)?,
),
});
}
Ok(out)
}
fn logical_columns(
base_name: &str,
array: LogicalArray,
split_matrix: bool,
) -> BuiltinResult<Vec<GroupColumn>> {
let rows = array.shape.first().copied().unwrap_or(array.data.len());
let cols = array.shape.get(1).copied().unwrap_or(1);
if !split_matrix || cols <= 1 || rows == 1 {
let len = array.data.len();
return Ok(vec![GroupColumn {
name: base_name.to_string(),
rows: len,
value: Value::LogicalArray(
LogicalArray::new(array.data, vec![len, 1]).map_err(grouping_error)?,
),
}]);
}
let mut out = Vec::with_capacity(cols);
for col in 0..cols {
let mut data = Vec::with_capacity(rows);
for row in 0..rows {
data.push(*array.data.get(row + col * rows).ok_or_else(|| {
grouping_error("grouping: logical array shape/data length mismatch")
})?);
}
out.push(GroupColumn {
name: format!("{base_name}{col_plus}", col_plus = col + 1),
rows,
value: Value::LogicalArray(
LogicalArray::new(data, vec![rows, 1]).map_err(grouping_error)?,
),
});
}
Ok(out)
}
fn string_columns(
base_name: &str,
array: StringArray,
split_matrix: bool,
) -> BuiltinResult<Vec<GroupColumn>> {
let rows = array.rows();
let cols = array.cols();
if !split_matrix || cols <= 1 || rows == 1 {
let len = array.data.len();
return Ok(vec![GroupColumn {
name: base_name.to_string(),
rows: len,
value: Value::StringArray(
StringArray::new(array.data, vec![len, 1]).map_err(grouping_error)?,
),
}]);
}
let mut out = Vec::with_capacity(cols);
for col in 0..cols {
let mut data = Vec::with_capacity(rows);
for row in 0..rows {
data.push(array.data[row + col * rows].clone());
}
out.push(GroupColumn {
name: format!("{base_name}{col_plus}", col_plus = col + 1),
rows,
value: Value::StringArray(
StringArray::new(data, vec![rows, 1]).map_err(grouping_error)?,
),
});
}
Ok(out)
}
fn cell_is_group_vector_list(cell: &CellArray) -> bool {
!cell.data.is_empty()
&& cell.data.iter().all(|value| {
matches!(
value,
Value::Tensor(_)
| Value::LogicalArray(_)
| Value::StringArray(_)
| Value::Object(_)
| Value::Cell(_)
)
})
}
fn build_grouping(columns: &[GroupColumn]) -> BuiltinResult<Grouping> {
build_grouping_with_options(
columns,
GroupOptions {
include_missing: false,
},
)
}
fn build_grouping_with_options(
columns: &[GroupColumn],
options: GroupOptions,
) -> BuiltinResult<Grouping> {
if columns.is_empty() {
return Err(grouping_error(
"grouping: expected at least one grouping variable",
));
}
let rows = columns[0].rows;
for column in columns {
if column.rows != rows {
return Err(grouping_error(format!(
"grouping: grouping variables must have matching row counts ({} vs {})",
rows, column.rows
)));
}
}
let mut buckets = BTreeMap::<Vec<Atom>, Vec<usize>>::new();
let mut row_keys = Vec::with_capacity(rows);
for row in 0..rows {
let key = columns
.iter()
.map(|column| atom_at(&column.value, row))
.collect::<BuiltinResult<Vec<_>>>()?;
if !options.include_missing && key.iter().any(|atom| matches!(atom, Atom::Missing)) {
row_keys.push(None);
continue;
}
buckets.entry(key.clone()).or_default().push(row);
row_keys.push(Some(key));
}
let keys = buckets.keys().cloned().collect::<Vec<_>>();
let mut key_to_index = BTreeMap::<Vec<Atom>, usize>::new();
let mut first_rows = Vec::<usize>::with_capacity(keys.len());
let mut row_groups = Vec::<Vec<usize>>::with_capacity(keys.len());
for (idx, key) in keys.iter().enumerate() {
key_to_index.insert(key.clone(), idx);
let rows = buckets
.get(key)
.cloned()
.expect("key collected from buckets must exist");
first_rows.push(*rows.first().expect("nonmissing group has rows"));
row_groups.push(rows);
}
let ids = row_keys
.into_iter()
.map(|key| {
key.and_then(|key| key_to_index.get(&key).copied())
.map(|idx| idx as f64 + 1.0)
.unwrap_or(f64::NAN)
})
.collect();
Ok(Grouping {
ids,
keys,
first_rows,
row_groups,
})
}
fn atom_at(value: &Value, row: usize) -> BuiltinResult<Atom> {
match value {
Value::Tensor(tensor) => {
let value = *tensor
.data
.get(row)
.ok_or_else(|| grouping_error("grouping: numeric row out of bounds"))?;
if value.is_nan() {
Ok(Atom::Missing)
} else {
Ok(Atom::Number(value))
}
}
Value::LogicalArray(array) => Ok(array
.data
.get(row)
.map(|flag| Atom::Logical(*flag != 0))
.unwrap_or(Atom::Missing)),
Value::StringArray(array) => Ok(array
.data
.get(row)
.map(|text| {
if is_missing_text(text) {
Atom::Missing
} else {
Atom::Text(text.clone())
}
})
.unwrap_or(Atom::Missing)),
Value::Cell(cell) => cell
.data
.get(row)
.map(scalar_atom)
.unwrap_or(Ok(Atom::Missing)),
Value::Object(object) if object.is_class("categorical") => {
let label = categorical_label_at(object, row);
Ok(match label.as_deref() {
None | Some("<undefined>") | Some("") => Atom::Missing,
Some(text) => Atom::Text(text.to_string()),
})
}
Value::Object(object) if object.is_class("datetime") => {
let serials = crate::builtins::datetime::serials_from_datetime_value(value)?;
let value = serials.data.get(row).copied().unwrap_or(f64::NAN);
if value.is_nan() {
Ok(Atom::Missing)
} else {
Ok(Atom::Number(value))
}
}
Value::Object(object) if object.is_class("duration") => {
let tensor = crate::builtins::duration::duration_tensor_from_duration_value(value)?;
let value = tensor.data.get(row).copied().unwrap_or(f64::NAN);
if value.is_nan() {
Ok(Atom::Missing)
} else {
Ok(Atom::Number(value))
}
}
other if row == 0 => scalar_atom(other),
_ => Ok(Atom::Missing),
}
}
fn scalar_atom(value: &Value) -> BuiltinResult<Atom> {
match value {
Value::Num(value) => {
if value.is_nan() {
Ok(Atom::Missing)
} else {
Ok(Atom::Number(*value))
}
}
Value::Int(value) => Ok(Atom::Number(value.to_f64())),
Value::Bool(flag) => Ok(Atom::Logical(*flag)),
Value::String(text) => {
if is_missing_text(text) {
Ok(Atom::Missing)
} else {
Ok(Atom::Text(text.clone()))
}
}
Value::CharArray(chars) if chars.rows == 1 => Ok(Atom::Text(chars.data.iter().collect())),
other => Ok(Atom::Text(format!("{other}"))),
}
}
fn findgroups_outputs(
columns: &[GroupColumn],
grouping: &Grouping,
table_mode: bool,
) -> BuiltinResult<Vec<Value>> {
let g = Value::Tensor(
Tensor::new(grouping.ids.clone(), vec![grouping.ids.len(), 1]).map_err(grouping_error)?,
);
let mut outputs = vec![g];
if table_mode {
let mut names = Vec::with_capacity(columns.len());
let mut values = Vec::with_capacity(columns.len());
for column in columns {
names.push(column.name.clone());
values.push(select_rows(&column.value, &grouping.first_rows)?);
}
outputs.push(table_from_columns(names, values)?);
} else {
outputs.extend(group_label_outputs(columns, grouping)?);
}
Ok(outputs)
}
fn group_label_outputs(columns: &[GroupColumn], grouping: &Grouping) -> BuiltinResult<Vec<Value>> {
columns
.iter()
.map(|column| select_rows(&column.value, &grouping.first_rows))
.collect()
}
fn split_option_tail(args: Vec<Value>) -> BuiltinResult<(Vec<Value>, Vec<Value>)> {
let mut option_start = args.len();
for (idx, value) in args.iter().enumerate() {
if is_option_name(value) {
option_start = idx;
break;
}
}
if option_start < args.len() && !(args.len() - option_start).is_multiple_of(2) {
return Err(grouping_error(
"groupcounts: name-value options must be provided in pairs",
));
}
Ok((args[..option_start].to_vec(), args[option_start..].to_vec()))
}
fn groupcounts_table(
object: ObjectInstance,
selector_args: Vec<Value>,
option_args: Vec<Value>,
) -> BuiltinResult<Value> {
let all_names = table_variable_names_from_object(&object)?;
if selector_args.len() > 1 {
return Err(grouping_error(
"groupcounts: table input accepts one grouping variable selector before options",
));
}
let selector = selector_args.first().ok_or_else(|| {
grouping_error("groupcounts: table input requires a grouping variable selector")
})?;
let options = GroupOptions::parse(&option_args, "groupcounts")?;
let selected = parse_name_selector(selector, &all_names, "groupcounts")?;
let variables = table_variables(&object)?;
let height = table_height(&object)?;
let mut columns = Vec::with_capacity(selected.len());
for name in &selected {
let value = variables
.fields
.get(name)
.cloned()
.ok_or_else(|| grouping_error(format!("groupcounts: missing variable '{name}'")))?;
columns.push(GroupColumn {
name: name.clone(),
rows: value_row_count(&value)?,
value,
});
}
let grouping = build_grouping_with_options(&columns, options)?;
let mut out_names = selected.clone();
let mut out_columns = group_label_outputs(&columns, &grouping)?;
out_names.push("GroupCount".to_string());
out_columns.push(Value::Tensor(
Tensor::new(
grouping
.row_groups
.iter()
.map(|rows| rows.len() as f64)
.collect(),
vec![grouping.keys.len(), 1],
)
.map_err(grouping_error)?,
));
out_names.push("Percent".to_string());
out_columns.push(Value::Tensor(
Tensor::new(
grouping
.row_groups
.iter()
.map(|rows| {
if height == 0 {
0.0
} else {
rows.len() as f64 * 100.0 / height as f64
}
})
.collect(),
vec![grouping.keys.len(), 1],
)
.map_err(grouping_error)?,
));
table_from_columns(out_names, out_columns)
}
async fn accumarray_impl(subs: Value, data: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let index_rows = accumarray_subscripts(subs)?;
let rows = index_rows.len();
let data_values = accumarray_data_values(data, rows)?;
let output_shape = if let Some(size_value) = rest.first() {
if is_empty_value(size_value) {
infer_accumarray_shape(&index_rows)
} else {
parse_size_vector(size_value, "accumarray")?
}
} else {
infer_accumarray_shape(&index_rows)
};
let fun = rest.get(1).filter(|value| !is_empty_value(value)).cloned();
let fill = rest.get(2).filter(|value| !is_empty_value(value)).cloned();
let issparse = rest
.get(3)
.map(|value| bool_scalar(value, "accumarray issparse"))
.transpose()?
.unwrap_or(false);
let output_len = checked_element_count(&output_shape, "accumarray")?;
if output_len > MAX_MATERIALIZED_ELEMENTS {
return Err(too_large_error("accumarray: output is too large"));
}
let mut buckets: BTreeMap<usize, Vec<f64>> = BTreeMap::new();
for (idx, subs) in index_rows.iter().enumerate() {
let lin = subscript_to_linear(subs, &output_shape)?;
buckets.entry(lin).or_default().push(data_values[idx]);
}
let fill_num = match fill.as_ref() {
None => 0.0,
Some(value) => numeric_scalar(value, "accumarray fill value")?,
};
let mut data_out = vec![fill_num; output_len];
if let Some(fun) = fun {
let mut cell_out = vec![Value::Num(fill_num); data_out.len()];
let mut all_numeric = true;
for (lin, values) in buckets {
let result = apply_accumarray_callback(fun.clone(), values).await?;
if let Some(num) = value_as_numeric_scalar(&result) {
data_out[lin] = num;
cell_out[lin] = Value::Num(num);
} else {
all_numeric = false;
cell_out[lin] = result;
}
}
if all_numeric {
return accumarray_numeric_output(data_out, output_shape, issparse);
}
if issparse {
return Err(grouping_error(
"accumarray: sparse output requires numeric scalar group results",
));
}
let (rows, cols) = shape_to_rows_cols(&output_shape)?;
return CellArray::new(cell_out, rows, cols)
.map(Value::Cell)
.map_err(grouping_error);
}
for (lin, values) in buckets {
data_out[lin] = values.iter().sum();
}
accumarray_numeric_output(data_out, output_shape, issparse)
}
fn accumarray_subscripts(subs: Value) -> BuiltinResult<Vec<Vec<usize>>> {
match subs {
Value::Tensor(tensor) => {
if tensor.data.is_empty() {
return Ok(Vec::new());
}
if tensor.cols() <= 1 || tensor.rows() == 1 {
tensor
.data
.into_iter()
.map(|value| Ok(vec![positive_integer(value, "accumarray subscript")?]))
.collect()
} else {
let mut out = Vec::with_capacity(tensor.rows());
for row in 0..tensor.rows() {
let mut subs = Vec::with_capacity(tensor.cols());
for col in 0..tensor.cols() {
subs.push(positive_integer(
tensor.get2(row, col).map_err(grouping_error)?,
"accumarray subscript",
)?);
}
out.push(subs);
}
Ok(out)
}
}
Value::Cell(cell) => {
let mut columns = Vec::with_capacity(cell.data.len());
for value in cell.data {
let column = accumarray_subscripts(value)?;
if column.iter().any(|row| row.len() != 1) {
return Err(grouping_error(
"accumarray: cell subscript entries must be vectors",
));
}
columns.push(column.into_iter().map(|row| row[0]).collect::<Vec<_>>());
}
let rows = columns.first().map(Vec::len).unwrap_or(0);
for column in &columns {
if column.len() != rows {
return Err(grouping_error(
"accumarray: cell subscript vectors must have equal length",
));
}
}
let mut out = Vec::with_capacity(rows);
for row in 0..rows {
out.push(columns.iter().map(|column| column[row]).collect());
}
Ok(out)
}
other => Err(grouping_error(format!(
"accumarray: unsupported subscript input {other:?}"
))),
}
}
fn accumarray_data_values(data: Value, rows: usize) -> BuiltinResult<Vec<f64>> {
match data {
Value::Num(value) => Ok(vec![value; rows]),
Value::Int(value) => Ok(vec![value.to_f64(); rows]),
Value::Bool(value) => Ok(vec![if value { 1.0 } else { 0.0 }; rows]),
Value::Tensor(tensor) => {
if tensor.data.len() == 1 {
Ok(vec![tensor.data[0]; rows])
} else if tensor.data.len() == rows {
Ok(tensor.data)
} else {
Err(grouping_error(
"accumarray: data must be scalar or match subscript row count",
))
}
}
Value::LogicalArray(array) => {
if array.data.len() == 1 {
Ok(vec![if array.data[0] != 0 { 1.0 } else { 0.0 }; rows])
} else if array.data.len() == rows {
Ok(array
.data
.into_iter()
.map(|flag| if flag != 0 { 1.0 } else { 0.0 })
.collect())
} else {
Err(grouping_error(
"accumarray: data must be scalar or match subscript row count",
))
}
}
other => Err(grouping_error(format!(
"accumarray: unsupported data input {other:?}"
))),
}
}
fn infer_accumarray_shape(index_rows: &[Vec<usize>]) -> Vec<usize> {
let dims = index_rows.first().map(Vec::len).unwrap_or(1).max(1);
let mut shape = vec![0usize; dims];
for row in index_rows {
for (dim, idx) in row.iter().enumerate() {
shape[dim] = shape[dim].max(*idx);
}
}
if dims == 1 {
vec![shape[0], 1]
} else {
shape
}
}
fn subscript_to_linear(subs: &[usize], shape: &[usize]) -> BuiltinResult<usize> {
if subs.len() > shape.len() {
return Err(grouping_error("accumarray: too many subscript dimensions"));
}
let mut linear = 0usize;
let mut stride = 1usize;
for (dim, &size) in shape.iter().enumerate() {
let sub = subs.get(dim).copied().unwrap_or(1);
if sub == 0 || sub > size {
return Err(grouping_error("accumarray: subscript exceeds output size"));
}
linear = linear
.checked_add(
(sub - 1)
.checked_mul(stride)
.ok_or_else(|| too_large_error("accumarray: output linear index overflow"))?,
)
.ok_or_else(|| too_large_error("accumarray: output linear index overflow"))?;
stride = stride
.checked_mul(size)
.ok_or_else(|| too_large_error("accumarray: output size overflow"))?;
}
Ok(linear)
}
async fn apply_accumarray_callback(func: Value, values: Vec<f64>) -> BuiltinResult<Value> {
if let Some(name) = function_name(&func) {
match name.to_ascii_lowercase().as_str() {
"sum" => return Ok(Value::Num(values.iter().sum())),
"mean" => {
return Ok(Value::Num(if values.is_empty() {
f64::NAN
} else {
values.iter().sum::<f64>() / values.len() as f64
}))
}
"numel" | "length" => return Ok(Value::Num(values.len() as f64)),
"min" => return Ok(Value::Num(values.into_iter().fold(f64::INFINITY, f64::min))),
"max" => {
return Ok(Value::Num(
values.into_iter().fold(f64::NEG_INFINITY, f64::max),
))
}
_ => {}
}
}
let arg =
Value::Tensor(Tensor::new(values.clone(), vec![values.len(), 1]).map_err(grouping_error)?);
call_feval_async_with_outputs(func, &[arg], 1)
.await
.map_err(|err| callback_error("accumarray: callback failed", Some(err)))
}
fn accumarray_numeric_output(
data: Vec<f64>,
shape: Vec<usize>,
issparse: bool,
) -> BuiltinResult<Value> {
if issparse {
let (rows, cols) = shape_to_rows_cols(&shape)?;
if shape.len() > 2 {
return Err(grouping_error(
"accumarray: sparse output is only supported for 2-D results",
));
}
let mut col_ptrs = Vec::with_capacity(cols + 1);
let mut row_indices = Vec::new();
let mut values = Vec::new();
col_ptrs.push(0);
for col in 0..cols {
for row in 0..rows {
let value = data[row + col * rows];
if value != 0.0 {
row_indices.push(row);
values.push(value);
}
}
col_ptrs.push(values.len());
}
return SparseTensor::new(rows, cols, col_ptrs, row_indices, values)
.map(Value::SparseTensor)
.map_err(grouping_error);
}
Tensor::new(data, shape)
.map(Value::Tensor)
.map_err(grouping_error)
}
fn discretize_impl(x: Value, edges_or_n: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let values = numeric_values(&x, "discretize X")?;
let shape = value_shape(&x);
let (edges, labels, included_right) = parse_discretize_args(&values, edges_or_n, rest)?;
if edges.len() < 2 {
return Err(grouping_error(
"discretize: at least two bin edges are required",
));
}
let bins = values
.iter()
.map(|value| discretize_one(*value, &edges, included_right))
.collect::<Vec<_>>();
if let Some(labels) = labels {
let data = bins
.iter()
.map(|bin| match bin {
Some(idx) => labels.get(*idx - 1).cloned().unwrap_or_default(),
None => String::new(),
})
.collect::<Vec<_>>();
return StringArray::new(data, shape)
.map(Value::StringArray)
.map_err(grouping_error);
}
Tensor::new(
bins.into_iter()
.map(|bin| bin.map(|idx| idx as f64).unwrap_or(f64::NAN))
.collect(),
shape,
)
.map(Value::Tensor)
.map_err(grouping_error)
}
fn parse_discretize_args(
values: &[f64],
edges_or_n: Value,
rest: Vec<Value>,
) -> BuiltinResult<(Vec<f64>, Option<Vec<String>>, bool)> {
let mut labels = None;
let mut included_right = false;
let mut idx = 0usize;
if let Some(first) = rest.first() {
if !is_option_name(first) {
labels = Some(string_list(first)?);
idx = 1;
}
}
while idx < rest.len() {
if idx + 1 >= rest.len() {
return Err(grouping_error(
"discretize: name-value options must be provided in pairs",
));
}
let name = scalar_text(&rest[idx], "discretize option")?;
if name.eq_ignore_ascii_case("IncludedEdge") {
let edge = scalar_text(&rest[idx + 1], "IncludedEdge")?;
included_right = match edge.to_ascii_lowercase().as_str() {
"right" => true,
"left" => false,
other => {
return Err(grouping_error(format!(
"discretize: unsupported IncludedEdge '{other}'"
)))
}
};
} else {
return Err(grouping_error(format!(
"discretize: unsupported option '{name}'"
)));
}
idx += 2;
}
let edges = match edges_or_n {
Value::Num(n) if is_positive_integer_f64(n) => equal_width_edges(values, n as usize)?,
Value::Int(n) if n.to_i64() > 0 => equal_width_edges(values, n.to_i64() as usize)?,
other => numeric_values(&other, "discretize edges")?,
};
if let Some(labels) = &labels {
if labels.len() != edges.len().saturating_sub(1) {
return Err(grouping_error(
"discretize: number of labels must match number of bins",
));
}
}
Ok((edges, labels, included_right))
}
fn discretize_one(value: f64, edges: &[f64], included_right: bool) -> Option<usize> {
if value.is_nan() {
return None;
}
for bin in 0..edges.len() - 1 {
let lower = edges[bin];
let upper = edges[bin + 1];
let hit = if included_right {
(value > lower || (bin == 0 && value == lower)) && value <= upper
} else {
value >= lower && (value < upper || (bin == edges.len() - 2 && value == upper))
};
if hit {
return Some(bin + 1);
}
}
None
}
fn equal_width_edges(values: &[f64], bins: usize) -> BuiltinResult<Vec<f64>> {
if bins == 0 {
return Err(grouping_error(
"discretize: number of bins must be positive",
));
}
let finite = values
.iter()
.copied()
.filter(|value| value.is_finite())
.collect::<Vec<_>>();
if finite.is_empty() {
return Err(grouping_error(
"discretize: cannot infer equal-width bins from all-missing data",
));
}
let min = finite.iter().copied().fold(f64::INFINITY, f64::min);
let max = finite.iter().copied().fold(f64::NEG_INFINITY, f64::max);
if min == max {
let half = 0.5;
return Ok((0..=bins)
.map(|idx| min - half + idx as f64 / bins as f64)
.collect());
}
let step = (max - min) / bins as f64;
Ok((0..=bins).map(|idx| min + idx as f64 * step).collect())
}
fn combinations_impl(first: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
let mut values = vec![first];
let mut options_start = rest.len();
let mut idx = 0usize;
while idx < rest.len() {
if is_option_name(&rest[idx]) {
options_start = idx;
break;
}
values.push(rest[idx].clone());
idx += 1;
}
let mut names = (0..values.len())
.map(|idx| format!("Var{}", idx + 1))
.collect::<Vec<_>>();
if options_start < rest.len() {
let mut opt = options_start;
while opt < rest.len() {
if opt + 1 >= rest.len() {
return Err(grouping_error(
"combinations: name-value options must be provided in pairs",
));
}
let name = scalar_text(&rest[opt], "combinations option")?;
if name.eq_ignore_ascii_case("VariableNames") {
names = string_list(&rest[opt + 1])?;
if names.len() != values.len() {
return Err(grouping_error(
"combinations: VariableNames length must match input count",
));
}
} else {
return Err(grouping_error(format!(
"combinations: unsupported option '{name}'"
)));
}
opt += 2;
}
}
let columns = values
.into_iter()
.map(vector_elements)
.collect::<BuiltinResult<Vec<_>>>()?;
let row_count = columns.iter().try_fold(1usize, |acc, column| {
acc.checked_mul(column.len())
.ok_or_else(|| too_large_error("combinations: output row count overflow"))
})?;
if row_count > MAX_MATERIALIZED_ELEMENTS {
return Err(too_large_error("combinations: output is too large"));
}
let mut out_columns = Vec::with_capacity(columns.len());
for col_idx in 0..columns.len() {
let repeat_inner = columns[col_idx + 1..]
.iter()
.map(Vec::len)
.product::<usize>()
.max(1);
let repeat_outer = columns[..col_idx]
.iter()
.map(Vec::len)
.product::<usize>()
.max(1);
let mut values = Vec::with_capacity(row_count);
for _ in 0..repeat_outer {
for item in &columns[col_idx] {
for _ in 0..repeat_inner {
values.push(item.clone());
}
}
}
out_columns.push(collect_column_values(values, row_count)?);
}
table_from_columns(names, out_columns)
}
fn parse_name_selector(
value: &Value,
names: &[String],
context: &str,
) -> BuiltinResult<Vec<String>> {
match value {
Value::String(text) => {
if names.contains(text) {
Ok(vec![text.clone()])
} else {
Err(grouping_error(format!(
"{context}: unknown variable '{text}'"
)))
}
}
Value::CharArray(chars) if chars.rows == 1 => {
let text: String = chars.data.iter().collect();
parse_name_selector(&Value::String(text), names, context)
}
Value::StringArray(array) => array
.data
.iter()
.map(|name| {
if names.contains(name) {
Ok(name.clone())
} else {
Err(grouping_error(format!(
"{context}: unknown variable '{name}'"
)))
}
})
.collect(),
Value::Cell(cell) => cell
.data
.iter()
.map(|value| scalar_text(value, context))
.map(|res| {
let name = res?;
if names.contains(&name) {
Ok(name)
} else {
Err(grouping_error(format!(
"{context}: unknown variable '{name}'"
)))
}
})
.collect(),
Value::Tensor(tensor) => tensor
.data
.iter()
.map(|value| {
let idx = positive_integer(*value, context)?;
names.get(idx - 1).cloned().ok_or_else(|| {
grouping_error(format!("{context}: variable index out of range"))
})
})
.collect(),
other => Err(grouping_error(format!(
"{context}: unsupported variable selector {other:?}"
))),
}
}
async fn gather_values(values: Vec<Value>) -> BuiltinResult<Vec<Value>> {
let mut out = Vec::with_capacity(values.len());
for value in values {
out.push(gather_if_needed_async(&value).await?);
}
Ok(out)
}
fn normalize_outputs(value: Value, requested: usize, context: &str) -> BuiltinResult<Vec<Value>> {
match value {
Value::OutputList(values) if values.len() == requested => Ok(values),
Value::OutputList(values) => Err(callback_error(
format!(
"{context}: callback returned {} outputs but {} were requested",
values.len(),
requested
),
None,
)),
value if requested == 1 => Ok(vec![value]),
_ => Err(callback_error(
format!("{context}: callback did not return the requested number of outputs"),
None,
)),
}
}
fn collect_group_results(values: Vec<Value>, rows: usize, context: &str) -> BuiltinResult<Value> {
if values
.iter()
.all(|value| value_as_numeric_scalar(value).is_some())
{
return Tensor::new(
values
.iter()
.map(|value| value_as_numeric_scalar(value).unwrap())
.collect(),
vec![rows, 1],
)
.map(Value::Tensor)
.map_err(grouping_error);
}
CellArray::new(values, rows, 1)
.map(Value::Cell)
.map_err(|err| callback_error(format!("{context}: {err}"), None))
}
fn collect_column_values(values: Vec<Value>, rows: usize) -> BuiltinResult<Value> {
if values
.iter()
.all(|value| value_as_numeric_scalar(value).is_some())
{
return Tensor::new(
values
.iter()
.map(|value| value_as_numeric_scalar(value).unwrap())
.collect(),
vec![rows, 1],
)
.map(Value::Tensor)
.map_err(grouping_error);
}
if values.iter().all(|value| matches!(value, Value::String(_))) {
return StringArray::new(
values
.into_iter()
.map(|value| match value {
Value::String(text) => text,
_ => unreachable!("checked above"),
})
.collect(),
vec![rows, 1],
)
.map(Value::StringArray)
.map_err(grouping_error);
}
if values.iter().all(|value| matches!(value, Value::Bool(_))) {
return LogicalArray::new(
values
.into_iter()
.map(|value| match value {
Value::Bool(flag) => u8::from(flag),
_ => unreachable!("checked above"),
})
.collect(),
vec![rows, 1],
)
.map(Value::LogicalArray)
.map_err(grouping_error);
}
CellArray::new(values, rows, 1)
.map(Value::Cell)
.map_err(grouping_error)
}
fn vector_elements(value: Value) -> BuiltinResult<Vec<Value>> {
match value {
Value::Tensor(tensor) => Ok(tensor.data.into_iter().map(Value::Num).collect()),
Value::LogicalArray(array) => Ok(array
.data
.into_iter()
.map(|flag| Value::Bool(flag != 0))
.collect()),
Value::StringArray(array) => Ok(array.data.into_iter().map(Value::String).collect()),
Value::Cell(cell) => Ok(cell.data),
Value::CharArray(chars) if chars.rows == 1 => Ok(chars
.data
.into_iter()
.map(|ch| Value::String(ch.to_string()))
.collect()),
Value::String(text) => Ok(vec![Value::String(text)]),
other => Ok(vec![other]),
}
}
fn multi_output(outputs: Vec<Value>) -> BuiltinResult<Value> {
if let Some(out_count) = crate::output_count::current_output_count() {
if out_count == 0 {
return Ok(Value::OutputList(Vec::new()));
}
return Ok(crate::output_count::output_list_with_padding(
out_count, outputs,
));
}
Ok(outputs
.into_iter()
.next()
.unwrap_or(Value::OutputList(Vec::new())))
}
fn numeric_values(value: &Value, context: &str) -> BuiltinResult<Vec<f64>> {
match value {
Value::Num(value) => Ok(vec![*value]),
Value::Int(value) => Ok(vec![value.to_f64()]),
Value::Bool(value) => Ok(vec![if *value { 1.0 } else { 0.0 }]),
Value::Tensor(tensor) => Ok(tensor.data.clone()),
Value::LogicalArray(array) => Ok(array
.data
.iter()
.map(|flag| if *flag != 0 { 1.0 } else { 0.0 })
.collect()),
Value::SparseTensor(sparse) => sparse
.to_dense()
.map(|tensor| tensor.data)
.map_err(grouping_error),
other => Err(grouping_error(format!(
"{context}: expected numeric input, got {other:?}"
))),
}
}
fn value_shape(value: &Value) -> Vec<usize> {
match value {
Value::Tensor(tensor) => tensor.shape.clone(),
Value::LogicalArray(array) => array.shape.clone(),
Value::StringArray(array) => array.shape.clone(),
Value::SparseTensor(sparse) => sparse.shape(),
Value::Num(_) | Value::Int(_) | Value::Bool(_) => vec![1, 1],
_ => vec![1, 1],
}
}
fn parse_size_vector(value: &Value, context: &str) -> BuiltinResult<Vec<usize>> {
let values = numeric_values(value, context)?;
let mut dims = Vec::with_capacity(values.len());
for value in values {
dims.push(nonnegative_integer(value, context)?);
}
if dims.is_empty() {
return Err(grouping_error(format!(
"{context}: size vector must not be empty"
)));
}
if dims.len() == 1 {
dims.push(1);
}
Ok(dims)
}
fn shape_to_rows_cols(shape: &[usize]) -> BuiltinResult<(usize, usize)> {
let rows = shape.first().copied().unwrap_or(0);
let cols = if shape.len() <= 1 {
1
} else if shape.len() == 2 {
shape[1]
} else {
shape[1..].iter().product()
};
Ok((rows, cols))
}
fn checked_element_count(shape: &[usize], context: &str) -> BuiltinResult<usize> {
shape.iter().try_fold(1usize, |acc, dim| {
acc.checked_mul(*dim)
.ok_or_else(|| too_large_error(format!("{context}: output size overflow")))
})
}
fn scalar_text(value: &Value, context: &str) -> BuiltinResult<String> {
match value {
Value::String(text) => Ok(text.clone()),
Value::CharArray(chars) if chars.rows == 1 => Ok(chars.data.iter().collect()),
Value::StringArray(array) if array.data.len() == 1 => Ok(array.data[0].clone()),
other => Err(grouping_error(format!(
"{context}: expected text scalar, got {other:?}"
))),
}
}
fn string_list(value: &Value) -> BuiltinResult<Vec<String>> {
match value {
Value::String(text) => Ok(vec![text.clone()]),
Value::CharArray(chars) if chars.rows == 1 => Ok(vec![chars.data.iter().collect()]),
Value::StringArray(array) => Ok(array.data.clone()),
Value::Cell(cell) => cell
.data
.iter()
.map(|value| scalar_text(value, "string list"))
.collect(),
other => Err(grouping_error(format!("expected text list, got {other:?}"))),
}
}
fn bool_scalar(value: &Value, context: &str) -> BuiltinResult<bool> {
match value {
Value::Bool(flag) => Ok(*flag),
Value::Num(value) => Ok(*value != 0.0),
Value::Int(value) => Ok(value.to_i64() != 0),
Value::LogicalArray(array) if array.data.len() == 1 => Ok(array.data[0] != 0),
other => Err(grouping_error(format!(
"{context}: expected logical scalar, got {other:?}"
))),
}
}
fn numeric_scalar(value: &Value, context: &str) -> BuiltinResult<f64> {
value_as_numeric_scalar(value)
.ok_or_else(|| grouping_error(format!("{context}: expected numeric scalar")))
}
fn value_as_numeric_scalar(value: &Value) -> Option<f64> {
match value {
Value::Num(value) => Some(*value),
Value::Int(value) => Some(value.to_f64()),
Value::Bool(value) => Some(if *value { 1.0 } else { 0.0 }),
Value::Tensor(tensor) if tensor.data.len() == 1 => Some(tensor.data[0]),
Value::LogicalArray(array) if array.data.len() == 1 => {
Some(if array.data[0] != 0 { 1.0 } else { 0.0 })
}
_ => None,
}
}
fn is_empty_value(value: &Value) -> bool {
match value {
Value::Tensor(tensor) => tensor.data.is_empty(),
Value::StringArray(array) => array.data.is_empty(),
Value::Cell(cell) => cell.data.is_empty(),
Value::CharArray(chars) => chars.data.is_empty(),
_ => false,
}
}
fn is_option_name(value: &Value) -> bool {
scalar_text(value, "option")
.map(|text| {
matches!(
text.to_ascii_lowercase().as_str(),
"includededge" | "variablenames" | "includemissinggroups" | "includeemptygroups"
)
})
.unwrap_or(false)
}
fn function_name(value: &Value) -> Option<&str> {
match value {
Value::FunctionHandle(name)
| Value::ExternalFunctionHandle(name)
| Value::MethodFunctionHandle(name)
| Value::BoundFunctionHandle { name, .. }
| Value::String(name) => Some(name.as_str()),
_ => None,
}
}
fn positive_integer(value: f64, context: &str) -> BuiltinResult<usize> {
if is_positive_integer_f64(value) {
Ok(value as usize)
} else {
Err(grouping_error(format!(
"{context}: expected positive integer"
)))
}
}
fn nonnegative_integer(value: f64, context: &str) -> BuiltinResult<usize> {
if value.is_finite() && value >= 0.0 && value.fract() == 0.0 {
Ok(value as usize)
} else {
Err(grouping_error(format!(
"{context}: expected nonnegative integer"
)))
}
}
fn is_positive_integer_f64(value: f64) -> bool {
value.is_finite() && value > 0.0 && value.fract() == 0.0
}
fn format_key_number(value: f64) -> String {
if value.fract() == 0.0 && value.abs() < 1e15 {
format!("{}", value as i64)
} else {
let mut text = format!("{value:.12}");
while text.contains('.') && text.ends_with('0') {
text.pop();
}
if text.ends_with('.') {
text.pop();
}
text
}
}
fn is_missing_text(text: &str) -> bool {
text.eq_ignore_ascii_case("<missing>")
}
#[cfg(test)]
mod tests {
use super::*;
use futures::executor::block_on;
#[test]
fn accumarray_sums_vector_and_matrix_subscripts() {
let out = block_on(accumarray_builtin(
Value::Tensor(Tensor::new(vec![1.0, 3.0, 4.0, 2.0, 4.0, 1.0], vec![6, 1]).unwrap()),
Value::Tensor(Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6, 1]).unwrap()),
Vec::new(),
))
.unwrap();
match out {
Value::Tensor(tensor) => assert_eq!(tensor.data, vec![7.0, 4.0, 2.0, 8.0]),
other => panic!("expected tensor, got {other:?}"),
}
let out = block_on(accumarray_builtin(
Value::Tensor(
Tensor::new(
vec![1.0, 2.0, 3.0, 1.0, 2.0, 4.0, 1.0, 2.0, 2.0, 1.0, 2.0, 1.0],
vec![6, 2],
)
.unwrap(),
),
Value::Tensor(Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![6, 1]).unwrap()),
Vec::new(),
))
.unwrap();
match out {
Value::Tensor(tensor) => {
assert_eq!(tensor.shape, vec![4, 2]);
assert_eq!(tensor.data, vec![5.0, 0.0, 0.0, 6.0, 0.0, 7.0, 3.0, 0.0]);
}
other => panic!("expected tensor, got {other:?}"),
}
}
#[test]
fn discretize_assigns_bins_and_labels() {
let out = block_on(discretize_builtin(
Value::Tensor(Tensor::new(vec![0.0, 0.2, 1.0, 2.5], vec![4, 1]).unwrap()),
Value::Tensor(Tensor::new(vec![0.0, 1.0, 2.0], vec![1, 3]).unwrap()),
Vec::new(),
))
.unwrap();
match out {
Value::Tensor(tensor) => {
assert_eq!(tensor.data[0], 1.0);
assert_eq!(tensor.data[1], 1.0);
assert_eq!(tensor.data[2], 2.0);
assert!(tensor.data[3].is_nan());
}
other => panic!("expected tensor, got {other:?}"),
}
}
#[test]
fn findgroups_groupcounts_and_grp2idx_share_order() {
let groups = Value::StringArray(
StringArray::new(
vec!["b".into(), "a".into(), "b".into(), "<missing>".into()],
vec![4, 1],
)
.unwrap(),
);
let out = block_on(findgroups_builtin(groups.clone(), Vec::new())).unwrap();
match out {
Value::Tensor(tensor) => {
assert_eq!(tensor.data[0], 2.0);
assert_eq!(tensor.data[1], 1.0);
assert_eq!(tensor.data[2], 2.0);
assert!(tensor.data[3].is_nan());
}
other => panic!("expected tensor, got {other:?}"),
}
let counted = block_on(groupcounts_builtin(groups.clone(), Vec::new())).unwrap();
match counted {
Value::Tensor(tensor) => assert_eq!(tensor.data, vec![1.0, 2.0]),
other => panic!("expected tensor, got {other:?}"),
}
let indexed = block_on(grp2idx_builtin(groups)).unwrap();
match indexed {
Value::Tensor(tensor) => assert!(tensor.data[3].is_nan()),
other => panic!("expected tensor, got {other:?}"),
}
let empty_is_group = block_on(groupcounts_builtin(
Value::StringArray(
StringArray::new(vec![String::new(), "a".into(), String::new()], vec![3, 1])
.unwrap(),
),
Vec::new(),
))
.unwrap();
match empty_is_group {
Value::Tensor(tensor) => assert_eq!(tensor.data, vec![2.0, 1.0]),
other => panic!("expected tensor, got {other:?}"),
}
}
#[test]
fn combinations_returns_table_columns() {
let out = block_on(combinations_builtin(
Value::Tensor(Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap()),
vec![Value::StringArray(
StringArray::new(vec!["x".into(), "y".into()], vec![1, 2]).unwrap(),
)],
))
.unwrap();
let Value::Object(object) = out else {
panic!("expected table");
};
assert!(is_tabular_object(&object));
assert_eq!(table_height(&object).unwrap(), 4);
}
#[test]
fn accumarray_supports_callbacks_and_sparse_output() {
let out = block_on(accumarray_builtin(
Value::Tensor(Tensor::new(vec![1.0, 1.0, 2.0, 2.0], vec![4, 1]).unwrap()),
Value::Tensor(Tensor::new(vec![1.0, 3.0, 5.0, 7.0], vec![4, 1]).unwrap()),
vec![
Value::Tensor(Tensor::new(Vec::new(), vec![0, 0]).unwrap()),
Value::FunctionHandle("mean".into()),
],
))
.unwrap();
match out {
Value::Tensor(tensor) => assert_eq!(tensor.data, vec![2.0, 6.0]),
other => panic!("expected tensor, got {other:?}"),
}
let sparse = block_on(accumarray_builtin(
Value::Tensor(Tensor::new(vec![1.0, 3.0], vec![2, 1]).unwrap()),
Value::Num(1.0),
vec![
Value::Tensor(Tensor::new(vec![4.0, 1.0], vec![1, 2]).unwrap()),
Value::Tensor(Tensor::new(Vec::new(), vec![0, 0]).unwrap()),
Value::Tensor(Tensor::new(Vec::new(), vec![0, 0]).unwrap()),
Value::Bool(true),
],
))
.unwrap();
match sparse {
Value::SparseTensor(st) => {
assert_eq!(st.rows, 4);
assert_eq!(st.cols, 1);
assert_eq!(st.nnz(), 2);
}
other => panic!("expected sparse tensor, got {other:?}"),
}
}
#[test]
fn splitapply_invokes_callback_by_group() {
let out = block_on(splitapply_builtin(
Value::FunctionHandle("sum".into()),
Value::Tensor(Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1]).unwrap()),
vec![Value::Tensor(
Tensor::new(vec![2.0, 1.0, 2.0, 1.0], vec![4, 1]).unwrap(),
)],
))
.unwrap();
match out {
Value::Tensor(tensor) => assert_eq!(tensor.data, vec![6.0, 4.0]),
other => panic!("expected tensor, got {other:?}"),
}
let err = block_on(splitapply_builtin(
Value::FunctionHandle("sum".into()),
Value::Tensor(Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap()),
vec![Value::Tensor(
Tensor::new(vec![1.0, 1.0], vec![2, 1]).unwrap(),
)],
))
.unwrap_err();
assert!(err.message().contains("must have 2 rows"));
}
#[test]
fn groupcounts_table_returns_count_and_percent_columns() {
let table = table_from_columns(
vec!["G".into(), "X".into()],
vec![
Value::StringArray(
StringArray::new(vec!["b".into(), "a".into(), "b".into()], vec![3, 1]).unwrap(),
),
Value::Tensor(Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap()),
],
)
.unwrap();
let out = block_on(groupcounts_builtin(table, vec![Value::from("G")])).unwrap();
let Value::Object(object) = out else {
panic!("expected table");
};
let names = table_variable_names_from_object(&object).unwrap();
assert_eq!(names, vec!["G", "GroupCount", "Percent"]);
assert_eq!(table_height(&object).unwrap(), 2);
let table = table_from_columns(
vec!["G".into()],
vec![Value::StringArray(
StringArray::new(vec!["a".into(), "<missing>".into()], vec![2, 1]).unwrap(),
)],
)
.unwrap();
let out = block_on(groupcounts_builtin(
table,
vec![
Value::from("G"),
Value::from("IncludeMissingGroups"),
Value::Bool(true),
],
))
.unwrap();
let Value::Object(object) = out else {
panic!("expected table");
};
assert_eq!(table_height(&object).unwrap(), 2);
}
}