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