Skip to main content

reifydb_function/text/
pad_left.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::{constraint::bytes::MaxBytes, container::utf8::Utf8Container, r#type::Type};
6
7use crate::{ScalarFunction, ScalarFunctionContext, error::ScalarFunctionError, propagate_options};
8
9pub struct TextPadLeft;
10
11impl TextPadLeft {
12	pub fn new() -> Self {
13		Self
14	}
15}
16
17impl ScalarFunction for TextPadLeft {
18	fn scalar(&self, ctx: ScalarFunctionContext) -> crate::error::ScalarFunctionResult<ColumnData> {
19		if let Some(result) = propagate_options(self, &ctx) {
20			return result;
21		}
22
23		let columns = ctx.columns;
24		let row_count = ctx.row_count;
25
26		if columns.len() != 3 {
27			return Err(ScalarFunctionError::ArityMismatch {
28				function: ctx.fragment.clone(),
29				expected: 3,
30				actual: columns.len(),
31			});
32		}
33
34		let str_col = columns.get(0).unwrap();
35		let len_col = columns.get(1).unwrap();
36		let pad_col = columns.get(2).unwrap();
37
38		let pad_data = match pad_col.data() {
39			ColumnData::Utf8 {
40				container,
41				..
42			} => container,
43			other => {
44				return Err(ScalarFunctionError::InvalidArgumentType {
45					function: ctx.fragment.clone(),
46					argument_index: 2,
47					expected: vec![Type::Utf8],
48					actual: other.get_type(),
49				});
50			}
51		};
52
53		match str_col.data() {
54			ColumnData::Utf8 {
55				container: str_container,
56				..
57			} => {
58				let mut result_data = Vec::with_capacity(row_count);
59
60				for i in 0..row_count {
61					if !str_container.is_defined(i) || !pad_data.is_defined(i) {
62						result_data.push(String::new());
63						continue;
64					}
65
66					let target_len = match len_col.data() {
67						ColumnData::Int1(c) => c.get(i).map(|&v| v as i64),
68						ColumnData::Int2(c) => c.get(i).map(|&v| v as i64),
69						ColumnData::Int4(c) => c.get(i).map(|&v| v as i64),
70						ColumnData::Int8(c) => c.get(i).copied(),
71						ColumnData::Uint1(c) => c.get(i).map(|&v| v as i64),
72						ColumnData::Uint2(c) => c.get(i).map(|&v| v as i64),
73						ColumnData::Uint4(c) => c.get(i).map(|&v| v as i64),
74						_ => {
75							return Err(ScalarFunctionError::InvalidArgumentType {
76								function: ctx.fragment.clone(),
77								argument_index: 1,
78								expected: vec![
79									Type::Int1,
80									Type::Int2,
81									Type::Int4,
82									Type::Int8,
83								],
84								actual: len_col.data().get_type(),
85							});
86						}
87					};
88
89					match target_len {
90						Some(n) if n >= 0 => {
91							let s = &str_container[i];
92							let pad_char = &pad_data[i];
93							let char_count = s.chars().count();
94							let target = n as usize;
95
96							if char_count >= target {
97								result_data.push(s.to_string());
98							} else {
99								let pad_chars: Vec<char> = pad_char.chars().collect();
100								if pad_chars.is_empty() {
101									result_data.push(s.to_string());
102								} else {
103									let needed = target - char_count;
104									let mut padded = String::with_capacity(
105										s.len() + needed
106											* pad_chars[0].len_utf8(),
107									);
108									for j in 0..needed {
109										padded.push(
110											pad_chars[j % pad_chars.len()]
111										);
112									}
113									padded.push_str(s);
114									result_data.push(padded);
115								}
116							}
117						}
118						Some(_) => {
119							result_data.push(String::new());
120						}
121						None => {
122							result_data.push(String::new());
123						}
124					}
125				}
126
127				Ok(ColumnData::Utf8 {
128					container: Utf8Container::new(result_data),
129					max_bytes: MaxBytes::MAX,
130				})
131			}
132			other => Err(ScalarFunctionError::InvalidArgumentType {
133				function: ctx.fragment.clone(),
134				argument_index: 0,
135				expected: vec![Type::Utf8],
136				actual: other.get_type(),
137			}),
138		}
139	}
140
141	fn return_type(&self, _input_types: &[Type]) -> Type {
142		Type::Utf8
143	}
144}