reifydb_routine/function/series/
mod.rs

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

reifydb_routine/function/series/
mod.rs
