reifydb_routine/function/text/
pad_right.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::{
6	util::bitvec::BitVec,
7	value::{constraint::bytes::MaxBytes, container::utf8::Utf8Container, r#type::Type},
8};
9
10use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
11
12pub struct TextPadRight {
13	info: RoutineInfo,
14}
15
16impl Default for TextPadRight {
17	fn default() -> Self {
18		Self::new()
19	}
20}
21
22impl TextPadRight {
23	pub fn new() -> Self {
24		Self {
25			info: RoutineInfo::new("text::pad_right"),
26		}
27	}
28}
29
30impl<'a> Routine<FunctionContext<'a>> for TextPadRight {
31	fn info(&self) -> &RoutineInfo {
32		&self.info
33	}
34
35	fn return_type(&self, _input_types: &[Type]) -> Type {
36		Type::Utf8
37	}
38
39	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
40		if args.len() != 3 {
41			return Err(RoutineError::FunctionArityMismatch {
42				function: ctx.fragment.clone(),
43				expected: 3,
44				actual: args.len(),
45			});
46		}
47
48		let str_col = &args[0];
49		let len_col = &args[1];
50		let pad_col = &args[2];
51
52		let (str_data, str_bv) = str_col.unwrap_option();
53		let (len_data, len_bv) = len_col.unwrap_option();
54		let (pad_data, pad_bv) = pad_col.unwrap_option();
55		let row_count = str_data.len();
56
57		let pad_container = match pad_data {
58			ColumnBuffer::Utf8 {
59				container,
60				..
61			} => container,
62			other => {
63				return Err(RoutineError::FunctionInvalidArgumentType {
64					function: ctx.fragment.clone(),
65					argument_index: 2,
66					expected: vec![Type::Utf8],
67					actual: other.get_type(),
68				});
69			}
70		};
71
72		match str_data {
73			ColumnBuffer::Utf8 {
74				container: str_container,
75				..
76			} => {
77				let mut result_data = Vec::with_capacity(row_count);
78
79				for i in 0..row_count {
80					if !str_container.is_defined(i) || !pad_container.is_defined(i) {
81						result_data.push(String::new());
82						continue;
83					}
84
85					let target_len = match len_data {
86						ColumnBuffer::Int1(c) => c.get(i).map(|&v| v as i64),
87						ColumnBuffer::Int2(c) => c.get(i).map(|&v| v as i64),
88						ColumnBuffer::Int4(c) => c.get(i).map(|&v| v as i64),
89						ColumnBuffer::Int8(c) => c.get(i).copied(),
90						ColumnBuffer::Uint1(c) => c.get(i).map(|&v| v as i64),
91						ColumnBuffer::Uint2(c) => c.get(i).map(|&v| v as i64),
92						ColumnBuffer::Uint4(c) => c.get(i).map(|&v| v as i64),
93						_ => {
94							return Err(RoutineError::FunctionInvalidArgumentType {
95								function: ctx.fragment.clone(),
96								argument_index: 1,
97								expected: vec![
98									Type::Int1,
99									Type::Int2,
100									Type::Int4,
101									Type::Int8,
102								],
103								actual: len_data.get_type(),
104							});
105						}
106					};
107
108					match target_len {
109						Some(n) if n >= 0 => {
110							let s = str_container.get(i).unwrap();
111							let pad_char = pad_container.get(i).unwrap();
112							let char_count = s.chars().count();
113							let target = n as usize;
114
115							if char_count >= target {
116								result_data.push(s.to_string());
117							} else {
118								let pad_chars: Vec<char> = pad_char.chars().collect();
119								if pad_chars.is_empty() {
120									result_data.push(s.to_string());
121								} else {
122									let needed = target - char_count;
123									let mut padded = String::with_capacity(
124										s.len() + needed
125											* pad_chars[0].len_utf8(),
126									);
127									padded.push_str(s);
128									for j in 0..needed {
129										padded.push(
130											pad_chars[j % pad_chars.len()]
131										);
132									}
133									result_data.push(padded);
134								}
135							}
136						}
137						Some(_) => {
138							result_data.push(String::new());
139						}
140						None => {
141							result_data.push(String::new());
142						}
143					}
144				}
145
146				let result_col_data = ColumnBuffer::Utf8 {
147					container: Utf8Container::new(result_data),
148					max_bytes: MaxBytes::MAX,
149				};
150
151				// Combine all three bitvecs
152				let mut combined_bv: Option<BitVec> = None;
153				for bv in [str_bv, len_bv, pad_bv].into_iter().flatten() {
154					combined_bv = Some(match combined_bv {
155						Some(existing) => existing.and(bv),
156						None => bv.clone(),
157					});
158				}
159
160				let final_data = match combined_bv {
161					Some(bv) => ColumnBuffer::Option {
162						inner: Box::new(result_col_data),
163						bitvec: bv,
164					},
165					None => result_col_data,
166				};
167				Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), final_data)]))
168			}
169			other => Err(RoutineError::FunctionInvalidArgumentType {
170				function: ctx.fragment.clone(),
171				argument_index: 0,
172				expected: vec![Type::Utf8],
173				actual: other.get_type(),
174			}),
175		}
176	}
177}
178
179impl Function for TextPadRight {
180	fn kinds(&self) -> &[FunctionKind] {
181		&[FunctionKind::Scalar]
182	}
183}
reifydb_routine/function/text/pad_right.rs

reifydb_routine/function/text/
pad_right.rs
