Skip to main content

reifydb_function/math/scalar/
lcm.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use 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}