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