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

reifydb_routine/function/text/
pad_right.rs
