reifydb_function/series/
mod.rs1use reifydb_core::value::column::{Column, columns::Columns, data::ColumnData};
5use reifydb_type::value::r#type::Type;
6
7use crate::{GeneratorContext, GeneratorFunction, ScalarFunction, ScalarFunctionContext, error::ScalarFunctionError};
8
9pub struct GenerateSeries;
10
11impl GenerateSeries {
12 pub fn new() -> Self {
13 Self {}
14 }
15}
16
17impl GeneratorFunction for GenerateSeries {
18 fn generate<'a>(&self, ctx: GeneratorContext<'a>) -> crate::error::GeneratorFunctionResult<Columns> {
19 let params = &ctx.params;
21
22 assert_eq!(params.len(), 2, "generate_series requires exactly 2 parameters: start and end");
23
24 let start_column = params.get(0).unwrap();
26 let start_value = match start_column.data() {
27 ColumnData::Int4(container) => container.get(0).copied().unwrap_or(1),
28 _ => panic!("start parameter must be an integer"),
29 };
30
31 let end_column = params.get(1).unwrap();
33 let end_value = match end_column.data() {
34 ColumnData::Int4(container) => container.get(0).copied().unwrap_or(10),
35 _ => panic!("end parameter must be an integer"),
36 };
37
38 let series: Vec<i32> = (start_value..=end_value).collect();
40 let series_column = Column::int4("value", series);
41
42 Ok(Columns::new(vec![series_column]))
43 }
44}
45
46pub struct Series;
47
48impl Series {
49 pub fn new() -> Self {
50 Self {}
51 }
52}
53
54fn extract_i32(data: &ColumnData, index: usize) -> Option<i32> {
55 match data {
56 ColumnData::Int1(c) => c.get(index).map(|&v| v as i32),
57 ColumnData::Int2(c) => c.get(index).map(|&v| v as i32),
58 ColumnData::Int4(c) => c.get(index).copied(),
59 ColumnData::Int8(c) => c.get(index).map(|&v| v as i32),
60 ColumnData::Uint1(c) => c.get(index).map(|&v| v as i32),
61 ColumnData::Uint2(c) => c.get(index).map(|&v| v as i32),
62 ColumnData::Uint4(c) => c.get(index).map(|&v| v as i32),
63 _ => None,
64 }
65}
66
67impl ScalarFunction for Series {
68 fn scalar(&self, ctx: ScalarFunctionContext) -> crate::error::ScalarFunctionResult<ColumnData> {
69 let columns = ctx.columns;
70
71 if columns.len() != 2 {
72 return Err(ScalarFunctionError::ArityMismatch {
73 function: ctx.fragment.clone(),
74 expected: 2,
75 actual: columns.len(),
76 });
77 }
78
79 let start_column = columns.get(0).unwrap();
80 let start_value = extract_i32(start_column.data(), 0).unwrap_or(1);
81
82 let end_column = columns.get(1).unwrap();
83 let end_value = extract_i32(end_column.data(), 0).unwrap_or(10);
84
85 let series: Vec<i32> = (start_value..=end_value).collect();
86 Ok(ColumnData::int4(series))
87 }
88
89 fn return_type(&self, _input_types: &[Type]) -> Type {
90 Type::Int4
91 }
92}