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