reifydb_routine/function/math/
gcd.rs1use 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 Gcd {
10 info: RoutineInfo,
11}
12
13impl Default for Gcd {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl Gcd {
20 pub fn new() -> Self {
21 Self {
22 info: RoutineInfo::new("math::gcd"),
23 }
24 }
25}
26
27fn numeric_to_i64(data: &ColumnBuffer, i: usize) -> Option<i64> {
28 match data {
29 ColumnBuffer::Int1(c) => c.get(i).map(|&v| v as i64),
30 ColumnBuffer::Int2(c) => c.get(i).map(|&v| v as i64),
31 ColumnBuffer::Int4(c) => c.get(i).map(|&v| v as i64),
32 ColumnBuffer::Int8(c) => c.get(i).copied(),
33 ColumnBuffer::Int16(c) => c.get(i).map(|&v| v as i64),
34 ColumnBuffer::Uint1(c) => c.get(i).map(|&v| v as i64),
35 ColumnBuffer::Uint2(c) => c.get(i).map(|&v| v as i64),
36 ColumnBuffer::Uint4(c) => c.get(i).map(|&v| v as i64),
37 ColumnBuffer::Uint8(c) => c.get(i).map(|&v| v as i64),
38 _ => None,
39 }
40}
41
42fn compute_gcd(mut a: i64, mut b: i64) -> i64 {
43 a = a.abs();
44 b = b.abs();
45 while b != 0 {
46 let t = b;
47 b = a % b;
48 a = t;
49 }
50 a
51}
52
53impl<'a> Routine<FunctionContext<'a>> for Gcd {
54 fn info(&self) -> &RoutineInfo {
55 &self.info
56 }
57
58 fn return_type(&self, _input_types: &[Type]) -> Type {
59 Type::Int8
60 }
61
62 fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
63 if args.len() != 2 {
64 return Err(RoutineError::FunctionArityMismatch {
65 function: ctx.fragment.clone(),
66 expected: 2,
67 actual: args.len(),
68 });
69 }
70
71 let a_col = &args[0];
72 let b_col = &args[1];
73
74 let (a_data, a_bitvec) = a_col.unwrap_option();
75 let (b_data, b_bitvec) = b_col.unwrap_option();
76 let row_count = a_data.len();
77
78 let expected_types = vec![
79 Type::Int1,
80 Type::Int2,
81 Type::Int4,
82 Type::Int8,
83 Type::Uint1,
84 Type::Uint2,
85 Type::Uint4,
86 Type::Uint8,
87 ];
88 if !a_data.get_type().is_number() {
89 return Err(RoutineError::FunctionInvalidArgumentType {
90 function: ctx.fragment.clone(),
91 argument_index: 0,
92 expected: expected_types,
93 actual: a_data.get_type(),
94 });
95 }
96 if !b_data.get_type().is_number() {
97 return Err(RoutineError::FunctionInvalidArgumentType {
98 function: ctx.fragment.clone(),
99 argument_index: 1,
100 expected: expected_types,
101 actual: b_data.get_type(),
102 });
103 }
104
105 let mut result = Vec::with_capacity(row_count);
106 let mut res_bitvec = Vec::with_capacity(row_count);
107
108 for i in 0..row_count {
109 match (numeric_to_i64(a_data, i), numeric_to_i64(b_data, i)) {
110 (Some(a), Some(b)) => {
111 result.push(compute_gcd(a, b));
112 res_bitvec.push(true);
113 }
114 _ => {
115 result.push(0);
116 res_bitvec.push(false);
117 }
118 }
119 }
120
121 let result_data = ColumnBuffer::int8_with_bitvec(result, res_bitvec);
122 let combined_bitvec = match (a_bitvec, b_bitvec) {
123 (Some(a), Some(b)) => Some(a.and(b)),
124 (Some(a), None) => Some(a.clone()),
125 (None, Some(b)) => Some(b.clone()),
126 (None, None) => None,
127 };
128
129 let final_data = if let Some(bv) = combined_bitvec {
130 ColumnBuffer::Option {
131 inner: Box::new(result_data),
132 bitvec: bv,
133 }
134 } else {
135 result_data
136 };
137
138 Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), final_data)]))
139 }
140}
141
142impl Function for Gcd {
143 fn kinds(&self) -> &[FunctionKind] {
144 &[FunctionKind::Scalar]
145 }
146}