reifydb_function/text/
pad_right.rs1use 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 TextPadRight;
10
11impl TextPadRight {
12 pub fn new() -> Self {
13 Self
14 }
15}
16
17impl ScalarFunction for TextPadRight {
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 padded.push_str(s);
109 for j in 0..needed {
110 padded.push(
111 pad_chars[j % pad_chars.len()]
112 );
113 }
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}