Skip to main content

reifydb_function/series/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
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::{
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		// Extract parameters: start and end
23		let params = &ctx.params;
24
25		assert_eq!(params.len(), 2, "generate_series requires exactly 2 parameters: start and end");
26
27		// Get start value
28		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		// Get end value
35		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		// Generate the series
42		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}