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