reifydb_function/math/scalar/
lcm.rs1use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::r#type::Type;
6
7use crate::{
8 ScalarFunction, ScalarFunctionContext,
9 error::{ScalarFunctionError, ScalarFunctionResult},
10 propagate_options,
11};
12
13pub struct Lcm;
14
15impl Lcm {
16 pub fn new() -> Self {
17 Self
18 }
19}
20
21fn numeric_to_i64(data: &ColumnData, i: usize) -> Option<i64> {
22 match data {
23 ColumnData::Int1(c) => c.get(i).map(|&v| v as i64),
24 ColumnData::Int2(c) => c.get(i).map(|&v| v as i64),
25 ColumnData::Int4(c) => c.get(i).map(|&v| v as i64),
26 ColumnData::Int8(c) => c.get(i).copied(),
27 ColumnData::Int16(c) => c.get(i).map(|&v| v as i64),
28 ColumnData::Uint1(c) => c.get(i).map(|&v| v as i64),
29 ColumnData::Uint2(c) => c.get(i).map(|&v| v as i64),
30 ColumnData::Uint4(c) => c.get(i).map(|&v| v as i64),
31 ColumnData::Uint8(c) => c.get(i).map(|&v| v as i64),
32 _ => None,
33 }
34}
35
36fn compute_gcd(mut a: i64, mut b: i64) -> i64 {
37 a = a.abs();
38 b = b.abs();
39 while b != 0 {
40 let t = b;
41 b = a % b;
42 a = t;
43 }
44 a
45}
46
47fn compute_lcm(a: i64, b: i64) -> i64 {
48 if a == 0 || b == 0 {
49 return 0;
50 }
51 (a.abs() / compute_gcd(a, b)) * b.abs()
52}
53
54impl ScalarFunction for Lcm {
55 fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
56 if let Some(result) = propagate_options(self, &ctx) {
57 return result;
58 }
59 let columns = ctx.columns;
60 let row_count = ctx.row_count;
61
62 if columns.len() != 2 {
63 return Err(ScalarFunctionError::ArityMismatch {
64 function: ctx.fragment.clone(),
65 expected: 2,
66 actual: columns.len(),
67 });
68 }
69
70 let a_col = columns.get(0).unwrap();
71 let b_col = columns.get(1).unwrap();
72
73 if !a_col.data().get_type().is_number() {
74 return Err(ScalarFunctionError::InvalidArgumentType {
75 function: ctx.fragment.clone(),
76 argument_index: 0,
77 expected: vec![
78 Type::Int1,
79 Type::Int2,
80 Type::Int4,
81 Type::Int8,
82 Type::Uint1,
83 Type::Uint2,
84 Type::Uint4,
85 Type::Uint8,
86 ],
87 actual: a_col.data().get_type(),
88 });
89 }
90
91 if !b_col.data().get_type().is_number() {
92 return Err(ScalarFunctionError::InvalidArgumentType {
93 function: ctx.fragment.clone(),
94 argument_index: 1,
95 expected: vec![
96 Type::Int1,
97 Type::Int2,
98 Type::Int4,
99 Type::Int8,
100 Type::Uint1,
101 Type::Uint2,
102 Type::Uint4,
103 Type::Uint8,
104 ],
105 actual: b_col.data().get_type(),
106 });
107 }
108
109 let mut result = Vec::with_capacity(row_count);
110 let mut bitvec = Vec::with_capacity(row_count);
111
112 for i in 0..row_count {
113 match (numeric_to_i64(a_col.data(), i), numeric_to_i64(b_col.data(), i)) {
114 (Some(a), Some(b)) => {
115 result.push(compute_lcm(a, b));
116 bitvec.push(true);
117 }
118 _ => {
119 result.push(0);
120 bitvec.push(false);
121 }
122 }
123 }
124
125 Ok(ColumnData::int8_with_bitvec(result, bitvec))
126 }
127
128 fn return_type(&self, _input_types: &[Type]) -> Type {
129 Type::Int8
130 }
131}