Skip to main content

reifydb_function/
registry.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use std::{collections::HashMap, ops::Deref, sync::Arc};
5
6use crate::{
7	AggregateFunction, GeneratorFunction, ScalarFunction, blob, clock, date, datetime, duration, flow, identity,
8	is, math, meta, series, subscription, text, time,
9};
10
11#[derive(Clone)]
12pub struct Functions(Arc<FunctionsInner>);
13
14impl Functions {
15	pub fn empty() -> Functions {
16		Functions::builder().build()
17	}
18
19	pub fn builder() -> FunctionsBuilder {
20		FunctionsBuilder(FunctionsInner {
21			scalars: HashMap::new(),
22			aggregates: HashMap::new(),
23			generators: HashMap::new(),
24		})
25	}
26
27	pub fn defaults() -> FunctionsBuilder {
28		Functions::builder()
29			.register_aggregate("math::sum", math::aggregate::sum::Sum::new)
30			.register_aggregate("math::min", math::aggregate::min::Min::new)
31			.register_aggregate("math::max", math::aggregate::max::Max::new)
32			.register_aggregate("math::avg", math::aggregate::avg::Avg::new)
33			.register_aggregate("math::count", math::aggregate::count::Count::new)
34			.register_scalar("flow_node::to_json", flow::to_json::FlowNodeToJson::new)
35			.register_scalar("clock::now", clock::now::Now::new)
36			.register_scalar("clock::set", clock::set::Set::new)
37			.register_scalar("clock::advance", clock::advance::Advance::new)
38			.register_scalar("blob::b58", blob::b58::BlobB58::new)
39			.register_scalar("blob::b64", blob::b64::BlobB64::new)
40			.register_scalar("blob::b64url", blob::b64url::BlobB64url::new)
41			.register_scalar("blob::hex", blob::hex::BlobHex::new)
42			.register_scalar("blob::utf8", blob::utf8::BlobUtf8::new)
43			.register_scalar("math::abs", math::scalar::abs::Abs::new)
44			.register_scalar("math::acos", math::scalar::acos::Acos::new)
45			.register_scalar("math::asin", math::scalar::asin::Asin::new)
46			.register_scalar("math::atan", math::scalar::atan::Atan::new)
47			.register_scalar("math::atan2", math::scalar::atan2::Atan2::new)
48			.register_scalar("math::avg", math::scalar::avg::Avg::new)
49			.register_scalar("math::ceil", math::scalar::ceil::Ceil::new)
50			.register_scalar("math::clamp", math::scalar::clamp::Clamp::new)
51			.register_scalar("math::cos", math::scalar::cos::Cos::new)
52			.register_scalar("math::e", math::scalar::euler::Euler::new)
53			.register_scalar("math::exp", math::scalar::exp::Exp::new)
54			.register_scalar("math::floor", math::scalar::floor::Floor::new)
55			.register_scalar("math::gcd", math::scalar::gcd::Gcd::new)
56			.register_scalar("math::lcm", math::scalar::lcm::Lcm::new)
57			.register_scalar("math::log", math::scalar::log::Log::new)
58			.register_scalar("math::log10", math::scalar::log10::Log10::new)
59			.register_scalar("math::log2", math::scalar::log2::Log2::new)
60			.register_scalar("math::max", math::scalar::max::Max::new)
61			.register_scalar("math::min", math::scalar::min::Min::new)
62			.register_scalar("math::mod", math::scalar::modulo::Modulo::new)
63			.register_scalar("math::pi", math::scalar::pi::Pi::new)
64			.register_scalar("math::power", math::scalar::power::Power::new)
65			.register_scalar("math::round", math::scalar::round::Round::new)
66			.register_scalar("math::sign", math::scalar::sign::Sign::new)
67			.register_scalar("math::sin", math::scalar::sin::Sin::new)
68			.register_scalar("math::sqrt", math::scalar::sqrt::Sqrt::new)
69			.register_scalar("math::tan", math::scalar::tan::Tan::new)
70			.register_scalar("math::truncate", math::scalar::truncate::Truncate::new)
71			.register_scalar("date::year", date::year::DateYear::new)
72			.register_scalar("date::month", date::month::DateMonth::new)
73			.register_scalar("date::day", date::day::DateDay::new)
74			.register_scalar("date::day_of_year", date::day_of_year::DateDayOfYear::new)
75			.register_scalar("date::day_of_week", date::day_of_week::DateDayOfWeek::new)
76			.register_scalar("date::quarter", date::quarter::DateQuarter::new)
77			.register_scalar("date::week", date::week::DateWeek::new)
78			.register_scalar("date::is_leap_year", date::is_leap_year::DateIsLeapYear::new)
79			.register_scalar("date::days_in_month", date::days_in_month::DateDaysInMonth::new)
80			.register_scalar("date::end_of_month", date::end_of_month::DateEndOfMonth::new)
81			.register_scalar("date::start_of_month", date::start_of_month::DateStartOfMonth::new)
82			.register_scalar("date::start_of_year", date::start_of_year::DateStartOfYear::new)
83			.register_scalar("date::new", date::new::DateNew::new)
84			.register_scalar("date::now", date::now::DateNow::new)
85			.register_scalar("date::add", date::add::DateAdd::new)
86			.register_scalar("date::subtract", date::subtract::DateSubtract::new)
87			.register_scalar("date::diff", date::diff::DateDiff::new)
88			.register_scalar("date::trunc", date::trunc::DateTrunc::new)
89			.register_scalar("date::age", date::age::DateAge::new)
90			.register_scalar("date::format", date::format::DateFormat::new)
91			.register_scalar("time::hour", time::hour::TimeHour::new)
92			.register_scalar("time::minute", time::minute::TimeMinute::new)
93			.register_scalar("time::second", time::second::TimeSecond::new)
94			.register_scalar("time::nanosecond", time::nanosecond::TimeNanosecond::new)
95			.register_scalar("time::new", time::new::TimeNew::new)
96			.register_scalar("time::now", time::now::TimeNow::new)
97			.register_scalar("time::add", time::add::TimeAdd::new)
98			.register_scalar("time::subtract", time::subtract::TimeSubtract::new)
99			.register_scalar("time::diff", time::diff::TimeDiff::new)
100			.register_scalar("time::trunc", time::trunc::TimeTrunc::new)
101			.register_scalar("time::age", time::age::TimeAge::new)
102			.register_scalar("time::format", time::format::TimeFormat::new)
103			.register_scalar("datetime::year", datetime::year::DateTimeYear::new)
104			.register_scalar("datetime::month", datetime::month::DateTimeMonth::new)
105			.register_scalar("datetime::day", datetime::day::DateTimeDay::new)
106			.register_scalar("datetime::hour", datetime::hour::DateTimeHour::new)
107			.register_scalar("datetime::minute", datetime::minute::DateTimeMinute::new)
108			.register_scalar("datetime::second", datetime::second::DateTimeSecond::new)
109			.register_scalar("datetime::nanosecond", datetime::nanosecond::DateTimeNanosecond::new)
110			.register_scalar("datetime::day_of_year", datetime::day_of_year::DateTimeDayOfYear::new)
111			.register_scalar("datetime::day_of_week", datetime::day_of_week::DateTimeDayOfWeek::new)
112			.register_scalar("datetime::quarter", datetime::quarter::DateTimeQuarter::new)
113			.register_scalar("datetime::week", datetime::week::DateTimeWeek::new)
114			.register_scalar("datetime::date", datetime::date::DateTimeDate::new)
115			.register_scalar("datetime::time", datetime::time::DateTimeTime::new)
116			.register_scalar("datetime::epoch", datetime::epoch::DateTimeEpoch::new)
117			.register_scalar("datetime::epoch_millis", datetime::epoch_millis::DateTimeEpochMillis::new)
118			.register_scalar("datetime::new", datetime::new::DateTimeNew::new)
119			.register_scalar("datetime::now", datetime::now::DateTimeNow::new)
120			.register_scalar("datetime::from_epoch", datetime::from_epoch::DateTimeFromEpoch::new)
121			.register_scalar(
122				"datetime::from_epoch_millis",
123				datetime::from_epoch_millis::DateTimeFromEpochMillis::new,
124			)
125			.register_scalar("datetime::add", datetime::add::DateTimeAdd::new)
126			.register_scalar("datetime::subtract", datetime::subtract::DateTimeSubtract::new)
127			.register_scalar("datetime::diff", datetime::diff::DateTimeDiff::new)
128			.register_scalar("datetime::trunc", datetime::trunc::DateTimeTrunc::new)
129			.register_scalar("datetime::age", datetime::age::DateTimeAge::new)
130			.register_scalar("datetime::format", datetime::format::DateTimeFormat::new)
131			.register_scalar("duration::years", duration::years::DurationYears::new)
132			.register_scalar("duration::months", duration::months::DurationMonths::new)
133			.register_scalar("duration::weeks", duration::weeks::DurationWeeks::new)
134			.register_scalar("duration::days", duration::days::DurationDays::new)
135			.register_scalar("duration::hours", duration::hours::DurationHours::new)
136			.register_scalar("duration::minutes", duration::minutes::DurationMinutes::new)
137			.register_scalar("duration::seconds", duration::seconds::DurationSeconds::new)
138			.register_scalar("duration::millis", duration::millis::DurationMillis::new)
139			.register_scalar("duration::year", duration::years::DurationYears::new)
140			.register_scalar("duration::month", duration::months::DurationMonths::new)
141			.register_scalar("duration::week", duration::weeks::DurationWeeks::new)
142			.register_scalar("duration::day", duration::days::DurationDays::new)
143			.register_scalar("duration::hour", duration::hours::DurationHours::new)
144			.register_scalar("duration::minute", duration::minutes::DurationMinutes::new)
145			.register_scalar("duration::second", duration::seconds::DurationSeconds::new)
146			.register_scalar("duration::get_months", duration::get_months::DurationGetMonths::new)
147			.register_scalar("duration::get_days", duration::get_days::DurationGetDays::new)
148			.register_scalar("duration::get_nanos", duration::get_nanos::DurationGetNanos::new)
149			.register_scalar("duration::add", duration::add::DurationAdd::new)
150			.register_scalar("duration::subtract", duration::subtract::DurationSubtract::new)
151			.register_scalar("duration::negate", duration::negate::DurationNegate::new)
152			.register_scalar("duration::scale", duration::scale::DurationScale::new)
153			.register_scalar("duration::trunc", duration::trunc::DurationTrunc::new)
154			.register_scalar("duration::format", duration::format::DurationFormat::new)
155			.register_scalar("text::ascii", text::ascii::TextAscii::new)
156			.register_scalar("text::char", text::char::TextChar::new)
157			.register_scalar("text::concat", text::concat::TextConcat::new)
158			.register_scalar("text::contains", text::contains::TextContains::new)
159			.register_scalar("text::count", text::count::TextCount::new)
160			.register_scalar("text::ends_with", text::ends_with::TextEndsWith::new)
161			.register_scalar("text::index_of", text::index_of::TextIndexOf::new)
162			.register_scalar("text::pad_left", text::pad_left::TextPadLeft::new)
163			.register_scalar("text::pad_right", text::pad_right::TextPadRight::new)
164			.register_scalar("text::repeat", text::repeat::TextRepeat::new)
165			.register_scalar("text::replace", text::replace::TextReplace::new)
166			.register_scalar("text::reverse", text::reverse::TextReverse::new)
167			.register_scalar("text::starts_with", text::starts_with::TextStartsWith::new)
168			.register_scalar("text::length", text::length::TextLength::new)
169			.register_scalar("text::trim", text::trim::TextTrim::new)
170			.register_scalar("text::trim_end", text::trim_end::TextTrimEnd::new)
171			.register_scalar("text::trim_start", text::trim_start::TextTrimStart::new)
172			.register_scalar("text::upper", text::upper::TextUpper::new)
173			.register_scalar("text::lower", text::lower::TextLower::new)
174			.register_scalar("text::substring", text::substring::TextSubstring::new)
175			.register_scalar("text::format_bytes", text::format_bytes::FormatBytes::new)
176			.register_scalar("text::format_bytes_si", text::format_bytes_si::FormatBytesSi::new)
177			.register_scalar("meta::type", meta::r#type::Type::new)
178			.register_scalar("identity::id", identity::id::Id::new)
179			.register_scalar("is::some", is::some::IsSome::new)
180			.register_scalar("is::none", is::none::IsNone::new)
181			.register_scalar("is::type", is::r#type::IsType::new)
182			.register_scalar("is::root", is::root::IsRoot::new)
183			.register_scalar("is::anonymous", is::anonymous::IsAnonymous::new)
184			.register_scalar("gen::series", series::Series::new)
185			.register_generator("generate_series", series::GenerateSeries::new)
186			.register_generator("inspect_subscription", subscription::inspect::InspectSubscription::new)
187	}
188}
189
190impl Deref for Functions {
191	type Target = FunctionsInner;
192
193	fn deref(&self) -> &Self::Target {
194		&self.0
195	}
196}
197
198#[derive(Clone)]
199pub struct FunctionsInner {
200	scalars: HashMap<String, Arc<dyn Fn() -> Box<dyn ScalarFunction> + Send + Sync>>,
201	aggregates: HashMap<String, Arc<dyn Fn() -> Box<dyn AggregateFunction> + Send + Sync>>,
202	generators: HashMap<String, Arc<dyn Fn() -> Box<dyn GeneratorFunction> + Send + Sync>>,
203}
204
205impl FunctionsInner {
206	pub fn get_aggregate(&self, name: &str) -> Option<Box<dyn AggregateFunction>> {
207		self.aggregates.get(name).map(|func| func())
208	}
209
210	pub fn get_scalar(&self, name: &str) -> Option<Box<dyn ScalarFunction>> {
211		self.scalars.get(name).map(|func| func())
212	}
213
214	pub fn get_generator(&self, name: &str) -> Option<Box<dyn GeneratorFunction>> {
215		self.generators.get(name).map(|func| func())
216	}
217
218	pub fn scalar_names(&self) -> Vec<&str> {
219		self.scalars.keys().map(|s| s.as_str()).collect()
220	}
221
222	pub fn aggregate_names(&self) -> Vec<&str> {
223		self.aggregates.keys().map(|s| s.as_str()).collect()
224	}
225
226	pub fn generator_names(&self) -> Vec<&str> {
227		self.generators.keys().map(|s| s.as_str()).collect()
228	}
229
230	pub fn get_scalar_factory(&self, name: &str) -> Option<Arc<dyn Fn() -> Box<dyn ScalarFunction> + Send + Sync>> {
231		self.scalars.get(name).cloned()
232	}
233
234	pub fn get_aggregate_factory(
235		&self,
236		name: &str,
237	) -> Option<Arc<dyn Fn() -> Box<dyn AggregateFunction> + Send + Sync>> {
238		self.aggregates.get(name).cloned()
239	}
240}
241
242pub struct FunctionsBuilder(FunctionsInner);
243
244impl FunctionsBuilder {
245	pub fn register_scalar<F, A>(mut self, name: &str, init: F) -> Self
246	where
247		F: Fn() -> A + Send + Sync + 'static,
248		A: ScalarFunction + 'static,
249	{
250		self.0.scalars.insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn ScalarFunction>));
251
252		self
253	}
254
255	pub fn register_aggregate<F, A>(mut self, name: &str, init: F) -> Self
256	where
257		F: Fn() -> A + Send + Sync + 'static,
258		A: AggregateFunction + 'static,
259	{
260		self.0.aggregates
261			.insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn AggregateFunction>));
262
263		self
264	}
265
266	pub fn register_generator<F, G>(mut self, name: &str, init: F) -> Self
267	where
268		F: Fn() -> G + Send + Sync + 'static,
269		G: GeneratorFunction + 'static,
270	{
271		self.0.generators
272			.insert(name.to_string(), Arc::new(move || Box::new(init()) as Box<dyn GeneratorFunction>));
273
274		self
275	}
276
277	pub fn build(self) -> Functions {
278		Functions(Arc::new(self.0))
279	}
280}