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		let start_column = args.first().ok_or_else(|| RoutineError::FunctionArityMismatch {
46			function: ctx.fragment.clone(),
47			expected: 2,
48			actual: args.len(),
49		})?;
50		let start_value = match start_column.data() {
51			ColumnBuffer::Int4(container) => container.get(0).copied().unwrap_or(1),
52			_ => {
53				return Err(RoutineError::FunctionExecutionFailed {
54					function: ctx.fragment.clone(),
55					reason: "start parameter must be an integer".to_string(),
56				});
57			}
58		};
59
60		let end_column = args.get(1).ok_or_else(|| RoutineError::FunctionArityMismatch {
61			function: ctx.fragment.clone(),
62			expected: 2,
63			actual: args.len(),
64		})?;
65		let end_value = match end_column.data() {
66			ColumnBuffer::Int4(container) => container.get(0).copied().unwrap_or(10),
67			_ => {
68				return Err(RoutineError::FunctionExecutionFailed {
69					function: ctx.fragment.clone(),
70					reason: "end parameter must be an integer".to_string(),
71				});
72			}
73		};
74
75		let series: Vec<i32> = (start_value..=end_value).collect();
76		let series_column = ColumnWithName::int4("value", series);
77
78		Ok(Columns::new(vec![series_column]))
79	}
80}
81
82impl Function for GenerateSeries {
83	fn kinds(&self) -> &[FunctionKind] {
84		&[FunctionKind::Generator]
85	}
86}
87
88pub struct Series {
89	info: RoutineInfo,
90}
91
92impl Default for Series {
93	fn default() -> Self {
94		Self::new()
95	}
96}
97
98impl Series {
99	pub fn new() -> Self {
100		Self {
101			info: RoutineInfo::new("gen::series"),
102		}
103	}
104}
105
106fn extract_i32(data: &ColumnBuffer, index: usize) -> Option<i32> {
107	match data {
108		ColumnBuffer::Int1(c) => c.get(index).map(|&v| v as i32),
109		ColumnBuffer::Int2(c) => c.get(index).map(|&v| v as i32),
110		ColumnBuffer::Int4(c) => c.get(index).copied(),
111		ColumnBuffer::Int8(c) => c.get(index).map(|&v| v as i32),
112		ColumnBuffer::Uint1(c) => c.get(index).map(|&v| v as i32),
113		ColumnBuffer::Uint2(c) => c.get(index).map(|&v| v as i32),
114		ColumnBuffer::Uint4(c) => c.get(index).map(|&v| v as i32),
115		_ => None,
116	}
117}
118
119impl<'a> Routine<FunctionContext<'a>> for Series {
120	fn info(&self) -> &RoutineInfo {
121		&self.info
122	}
123
124	fn return_type(&self, _input_types: &[Type]) -> Type {
125		Type::Int4
126	}
127
128	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
129		if args.len() != 2 {
130			return Err(RoutineError::FunctionArityMismatch {
131				function: ctx.fragment.clone(),
132				expected: 2,
133				actual: args.len(),
134			});
135		}
136
137		let start_column = args.first().ok_or_else(|| RoutineError::FunctionArityMismatch {
138			function: ctx.fragment.clone(),
139			expected: 2,
140			actual: args.len(),
141		})?;
142		let start_value = extract_i32(start_column.data(), 0).unwrap_or(1);
143
144		let end_column = args.get(1).ok_or_else(|| RoutineError::FunctionArityMismatch {
145			function: ctx.fragment.clone(),
146			expected: 2,
147			actual: args.len(),
148		})?;
149		let end_value = extract_i32(end_column.data(), 0).unwrap_or(10);
150
151		let series: Vec<i32> = (start_value..=end_value).collect();
152		Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), ColumnBuffer::int4(series))]))
153	}
154}
155
156impl Function for Series {
157	fn kinds(&self) -> &[FunctionKind] {
158		&[FunctionKind::Scalar]
159	}
160}
reifydb_routine/function/series/mod.rs

reifydb_routine/function/series/
mod.rs
