use crate::core::{Result, Value};
use crate::functions::{
FunctionDataType, FunctionInfo, FunctionSignature, FunctionType, WindowFunction,
};
#[derive(Default)]
pub struct RowNumberFunction;
impl WindowFunction for RowNumberFunction {
fn name(&self) -> &str {
"ROW_NUMBER"
}
fn info(&self) -> FunctionInfo {
FunctionInfo::new(
"ROW_NUMBER",
FunctionType::Window,
"Returns the sequential row number within the current partition",
FunctionSignature::new(FunctionDataType::Integer, vec![], 0, 0),
)
}
fn process(
&self,
_partition: &[Value],
_order_by: &[Value],
current_row: usize,
) -> Result<Value> {
Ok(Value::Integer((current_row + 1) as i64))
}
fn clone_box(&self) -> Box<dyn WindowFunction> {
Box::new(RowNumberFunction)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_row_number_basic() {
let f = RowNumberFunction;
let partition = vec![];
let order_by = vec![];
assert_eq!(
f.process(&partition, &order_by, 0).unwrap(),
Value::Integer(1)
);
assert_eq!(
f.process(&partition, &order_by, 1).unwrap(),
Value::Integer(2)
);
assert_eq!(
f.process(&partition, &order_by, 9).unwrap(),
Value::Integer(10)
);
}
#[test]
fn test_row_number_with_partition() {
let f = RowNumberFunction;
let partition = vec![
Value::Integer(100),
Value::Integer(200),
Value::Integer(300),
];
let order_by = vec![];
assert_eq!(
f.process(&partition, &order_by, 0).unwrap(),
Value::Integer(1)
);
assert_eq!(
f.process(&partition, &order_by, 1).unwrap(),
Value::Integer(2)
);
assert_eq!(
f.process(&partition, &order_by, 2).unwrap(),
Value::Integer(3)
);
}
#[test]
fn test_row_number_info() {
let f = RowNumberFunction;
let info = f.info();
assert_eq!(info.name(), "ROW_NUMBER");
assert_eq!(info.function_type(), FunctionType::Window);
}
}