use crate::data::datatable::DataValue;
use crate::sql::functions::{ArgCount, FunctionCategory, FunctionSignature, SqlFunction};
use anyhow::Result;
use lazy_static::lazy_static;
use std::collections::HashMap;
use std::sync::Mutex;
lazy_static! {
static ref GROUP_NUM_MEMO: Mutex<HashMap<String, HashMap<String, i64>>> =
Mutex::new(HashMap::new());
}
pub struct GroupNumFunction;
impl GroupNumFunction {
pub fn new() -> Self {
Self
}
pub fn clear_memoization() {
let mut memo = GROUP_NUM_MEMO.lock().unwrap();
memo.clear();
}
}
impl SqlFunction for GroupNumFunction {
fn signature(&self) -> FunctionSignature {
FunctionSignature {
name: "GROUP_NUM",
category: FunctionCategory::Aggregate,
arg_count: ArgCount::Fixed(1),
description: "Assigns unique sequential numbers (starting from 0) to distinct values",
returns: "Integer - unique number for each distinct value",
examples: vec![
"SELECT order_id, GROUP_NUM(order_id) as grp_num FROM orders",
"SELECT customer, GROUP_NUM(customer) as cust_num FROM sales",
],
}
}
fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
if args.len() != 1 {
anyhow::bail!("GROUP_NUM requires exactly 1 argument");
}
let value_str = match &args[0] {
DataValue::Null => return Ok(DataValue::Null),
DataValue::String(s) => s.clone(),
DataValue::InternedString(s) => s.to_string(),
DataValue::Integer(i) => i.to_string(),
DataValue::Float(f) => f.to_string(),
DataValue::Boolean(b) => b.to_string(),
DataValue::DateTime(dt) => dt.to_string(),
DataValue::Vector(v) => {
let components: Vec<String> = v.iter().map(|f| f.to_string()).collect();
format!("[{}]", components.join(","))
}
};
let column_id = "_default_";
let mut memo = GROUP_NUM_MEMO.lock().unwrap();
let column_map = memo
.entry(column_id.to_string())
.or_insert_with(HashMap::new);
if let Some(&num) = column_map.get(&value_str) {
Ok(DataValue::Integer(num))
} else {
let new_num = column_map.len() as i64;
column_map.insert(value_str, new_num);
Ok(DataValue::Integer(new_num))
}
}
}
pub struct GroupNumWithContext;
impl GroupNumWithContext {
pub fn new() -> Self {
Self
}
pub fn evaluate_with_context(&self, value: &DataValue, column_name: &str) -> Result<DataValue> {
if matches!(value, DataValue::Null) {
return Ok(DataValue::Null);
}
let value_str = match value {
DataValue::String(s) => s.clone(),
DataValue::InternedString(s) => s.to_string(),
DataValue::Integer(i) => i.to_string(),
DataValue::Float(f) => f.to_string(),
DataValue::Boolean(b) => b.to_string(),
DataValue::DateTime(dt) => dt.to_string(),
DataValue::Vector(v) => {
let components: Vec<String> = v.iter().map(|f| f.to_string()).collect();
format!("[{}]", components.join(","))
}
DataValue::Null => unreachable!(),
};
let mut memo = GROUP_NUM_MEMO.lock().unwrap();
let column_map = memo
.entry(column_name.to_string())
.or_insert_with(HashMap::new);
if let Some(&num) = column_map.get(&value_str) {
Ok(DataValue::Integer(num))
} else {
let new_num = column_map.len() as i64;
column_map.insert(value_str, new_num);
Ok(DataValue::Integer(new_num))
}
}
pub fn clear_column(&self, column_name: &str) {
let mut memo = GROUP_NUM_MEMO.lock().unwrap();
memo.remove(column_name);
}
pub fn clear_all(&self) {
let mut memo = GROUP_NUM_MEMO.lock().unwrap();
memo.clear();
}
}