reifydb_routine/function/series/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::value::r#type::Type;
6
7use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
8
9pub struct GenerateSeries {
10	info: RoutineInfo,
11}
12
13impl Default for GenerateSeries {
14	fn default() -> Self {
15		Self::new()
16	}
17}
18
19impl GenerateSeries {
20	pub fn new() -> Self {
21		Self {
22			info: RoutineInfo::new("series::generate"),
23		}
24	}
25}
26
27impl<'a> Routine<FunctionContext<'a>> for GenerateSeries {
28	fn info(&self) -> &RoutineInfo {
29		&self.info
30	}
31
32	fn return_type(&self, _input_types: &[Type]) -> Type {
33		Type::Any
34	}
35
36	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
37		if args.len() != 2 {
38			return Err(RoutineError::FunctionArityMismatch {
39				function: ctx.fragment.clone(),
40				expected: 2,
41				actual: args.len(),
42			});
43		}
44
45		// Get start value
46		let start_column = args.first().ok_or_else(|| RoutineError::FunctionArityMismatch {
47			function: ctx.fragment.clone(),
48			expected: 2,
49			actual: args.len(),
50		})?;
51		let start_value = match start_column.data() {
52			ColumnBuffer::Int4(container) => container.get(0).copied().unwrap_or(1),
53			_ => {
54				return Err(RoutineError::FunctionExecutionFailed {
55					function: ctx.fragment.clone(),
56					reason: "start parameter must be an integer".to_string(),
57				});
58			}
59		};
60
61		// Get end value
62		let end_column = args.get(1).ok_or_else(|| RoutineError::FunctionArityMismatch {
63			function: ctx.fragment.clone(),
64			expected: 2,
65			actual: args.len(),
66		})?;
67		let end_value = match end_column.data() {
68			ColumnBuffer::Int4(container) => container.get(0).copied().unwrap_or(10),
69			_ => {
70				return Err(RoutineError::FunctionExecutionFailed {
71					function: ctx.fragment.clone(),
72					reason: "end parameter must be an integer".to_string(),
73				});
74			}
75		};
76
77		// Generate the series
78		let series: Vec<i32> = (start_value..=end_value).collect();
79		let series_column = ColumnWithName::int4("value", series);
80
81		Ok(Columns::new(vec![series_column]))
82	}
83}
84
85impl Function for GenerateSeries {
86	fn kinds(&self) -> &[FunctionKind] {
87		&[FunctionKind::Generator]
88	}
89}
90
91pub struct Series {
92	info: RoutineInfo,
93}
94
95impl Default for Series {
96	fn default() -> Self {
97		Self::new()
98	}
99}
100
101impl Series {
102	pub fn new() -> Self {
103		Self {
104			info: RoutineInfo::new("gen::series"),
105		}
106	}
107}
108
109fn extract_i32(data: &ColumnBuffer, index: usize) -> Option<i32> {
110	match data {
111		ColumnBuffer::Int1(c) => c.get(index).map(|&v| v as i32),
112		ColumnBuffer::Int2(c) => c.get(index).map(|&v| v as i32),
113		ColumnBuffer::Int4(c) => c.get(index).copied(),
114		ColumnBuffer::Int8(c) => c.get(index).map(|&v| v as i32),
115		ColumnBuffer::Uint1(c) => c.get(index).map(|&v| v as i32),
116		ColumnBuffer::Uint2(c) => c.get(index).map(|&v| v as i32),
117		ColumnBuffer::Uint4(c) => c.get(index).map(|&v| v as i32),
118		_ => None,
119	}
120}
121
122impl<'a> Routine<FunctionContext<'a>> for Series {
123	fn info(&self) -> &RoutineInfo {
124		&self.info
125	}
126
127	fn return_type(&self, _input_types: &[Type]) -> Type {
128		Type::Int4
129	}
130
131	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
132		if args.len() != 2 {
133			return Err(RoutineError::FunctionArityMismatch {
134				function: ctx.fragment.clone(),
135				expected: 2,
136				actual: args.len(),
137			});
138		}
139
140		let start_column = args.first().ok_or_else(|| RoutineError::FunctionArityMismatch {
141			function: ctx.fragment.clone(),
142			expected: 2,
143			actual: args.len(),
144		})?;
145		let start_value = extract_i32(start_column.data(), 0).unwrap_or(1);
146
147		let end_column = args.get(1).ok_or_else(|| RoutineError::FunctionArityMismatch {
148			function: ctx.fragment.clone(),
149			expected: 2,
150			actual: args.len(),
151		})?;
152		let end_value = extract_i32(end_column.data(), 0).unwrap_or(10);
153
154		let series: Vec<i32> = (start_value..=end_value).collect();
155		Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), ColumnBuffer::int4(series))]))
156	}
157}
158
159impl Function for Series {
160	fn kinds(&self) -> &[FunctionKind] {
161		&[FunctionKind::Scalar]
162	}
163}
reifydb_routine/function/series/mod.rs

reifydb_routine/function/series/
mod.rs
