use crate::data::datatable::DataValue;
use anyhow::Result;
#[derive(Debug, Clone)]
pub struct ExpansionResult {
pub values: Vec<DataValue>,
}
impl ExpansionResult {
pub fn new(values: Vec<DataValue>) -> Self {
Self { values }
}
pub fn row_count(&self) -> usize {
self.values.len()
}
}
pub trait RowExpander: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
fn expand(&self, value: &DataValue, args: &[DataValue]) -> Result<ExpansionResult>;
}
pub struct RowExpanderRegistry {
expanders: std::collections::HashMap<String, Box<dyn RowExpander>>,
}
impl RowExpanderRegistry {
pub fn new() -> Self {
let mut registry = Self {
expanders: std::collections::HashMap::new(),
};
registry.register(Box::new(unnest::UnnestExpander));
registry
}
pub fn register(&mut self, expander: Box<dyn RowExpander>) {
self.expanders
.insert(expander.name().to_uppercase(), expander);
}
pub fn get(&self, name: &str) -> Option<&dyn RowExpander> {
self.expanders.get(&name.to_uppercase()).map(|e| e.as_ref())
}
pub fn contains(&self, name: &str) -> bool {
self.expanders.contains_key(&name.to_uppercase())
}
pub fn list(&self) -> Vec<&str> {
self.expanders.keys().map(|s| s.as_str()).collect()
}
}
impl Default for RowExpanderRegistry {
fn default() -> Self {
Self::new()
}
}
pub mod unnest {
use super::*;
pub struct UnnestExpander;
impl RowExpander for UnnestExpander {
fn name(&self) -> &str {
"UNNEST"
}
fn description(&self) -> &str {
"Split a delimited string into multiple rows"
}
fn expand(&self, value: &DataValue, args: &[DataValue]) -> Result<ExpansionResult> {
let delimiter = match args.first() {
Some(DataValue::String(s)) => s.as_str(),
Some(_) => {
return Err(anyhow::anyhow!(
"UNNEST delimiter must be a string, got {:?}",
args.first()
))
}
None => return Err(anyhow::anyhow!("UNNEST requires a delimiter argument")),
};
let text = match value {
DataValue::String(s) => s.clone(),
DataValue::Null => {
return Ok(ExpansionResult::new(vec![DataValue::Null]));
}
other => other.to_string(),
};
let parts: Vec<DataValue> = if delimiter.is_empty() {
text.chars()
.map(|ch| DataValue::String(ch.to_string()))
.collect()
} else {
text.split(delimiter)
.filter(|s| !s.is_empty())
.map(|s| DataValue::String(s.to_string()))
.collect()
};
if parts.is_empty() {
Ok(ExpansionResult::new(vec![DataValue::Null]))
} else {
Ok(ExpansionResult::new(parts))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_unnest_basic() {
let expander = unnest::UnnestExpander;
let value = DataValue::String("A|B|C".to_string());
let delimiter = DataValue::String("|".to_string());
let result = expander.expand(&value, &[delimiter]).unwrap();
assert_eq!(result.row_count(), 3);
assert_eq!(result.values[0], DataValue::String("A".to_string()));
assert_eq!(result.values[1], DataValue::String("B".to_string()));
assert_eq!(result.values[2], DataValue::String("C".to_string()));
}
#[test]
fn test_unnest_null() {
let expander = unnest::UnnestExpander;
let value = DataValue::Null;
let delimiter = DataValue::String("|".to_string());
let result = expander.expand(&value, &[delimiter]).unwrap();
assert_eq!(result.row_count(), 1);
assert_eq!(result.values[0], DataValue::Null);
}
#[test]
fn test_registry() {
let registry = RowExpanderRegistry::new();
assert!(registry.contains("UNNEST"));
assert!(registry.contains("unnest"));
let expander = registry.get("UNNEST").unwrap();
assert_eq!(expander.name(), "UNNEST");
}
}