sql_cli/sql/functions/
string_utils.rs

1use crate::data::datatable::DataValue;
2use crate::sql::functions::{ArgCount, FunctionCategory, FunctionSignature, SqlFunction};
3use anyhow::Result;
4
5/// REPEAT function - Repeat a string n times
6pub struct RepeatFunction;
7
8impl SqlFunction for RepeatFunction {
9    fn signature(&self) -> FunctionSignature {
10        FunctionSignature {
11            name: "REPEAT",
12            category: FunctionCategory::String,
13            arg_count: ArgCount::Fixed(2),
14            description: "Repeat a string n times",
15            returns: "String containing the input repeated n times",
16            examples: vec![
17                "SELECT REPEAT('*', 5)  -- Returns '*****'",
18                "SELECT REPEAT('ab', 3)  -- Returns 'ababab'",
19                "SELECT REPEAT('=', COUNT(*) / 10) FROM table  -- Create histogram",
20            ],
21        }
22    }
23
24    fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
25        if args.len() != 2 {
26            return Err(anyhow::anyhow!("REPEAT expects exactly 2 arguments"));
27        }
28
29        let string = match &args[0] {
30            DataValue::String(s) => s.clone(),
31            DataValue::Null => return Ok(DataValue::Null),
32            _ => args[0].to_string(),
33        };
34
35        let count = match &args[1] {
36            DataValue::Integer(n) => *n,
37            DataValue::Float(f) => *f as i64,
38            DataValue::Null => return Ok(DataValue::Null),
39            _ => {
40                return Err(anyhow::anyhow!(
41                    "REPEAT count must be a number, got {:?}",
42                    args[1]
43                ))
44            }
45        };
46
47        if count < 0 {
48            return Err(anyhow::anyhow!("REPEAT count cannot be negative"));
49        }
50
51        if count == 0 {
52            return Ok(DataValue::String(String::new()));
53        }
54
55        // Limit repetitions to prevent memory issues
56        if count > 10000 {
57            return Err(anyhow::anyhow!("REPEAT count too large (max 10000)"));
58        }
59
60        Ok(DataValue::String(string.repeat(count as usize)))
61    }
62}
63
64/// LPAD function - Left pad a string to a certain length
65pub struct LPadFunction;
66
67impl SqlFunction for LPadFunction {
68    fn signature(&self) -> FunctionSignature {
69        FunctionSignature {
70            name: "LPAD",
71            category: FunctionCategory::String,
72            arg_count: ArgCount::Range(2, 3),
73            description: "Left pad a string to a certain length",
74            returns: "String padded on the left to specified length",
75            examples: vec![
76                "SELECT LPAD('5', 3, '0')  -- Returns '005'",
77                "SELECT LPAD('hello', 10, ' ')  -- Returns '     hello'",
78                "SELECT LPAD('abc', 5)  -- Returns '  abc' (default pad is space)",
79            ],
80        }
81    }
82
83    fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
84        if args.len() < 2 || args.len() > 3 {
85            return Err(anyhow::anyhow!("LPAD expects 2 or 3 arguments"));
86        }
87
88        let string = match &args[0] {
89            DataValue::String(s) => s.clone(),
90            DataValue::Null => return Ok(DataValue::Null),
91            _ => args[0].to_string(),
92        };
93
94        let length = match &args[1] {
95            DataValue::Integer(n) => *n as usize,
96            DataValue::Float(f) => *f as usize,
97            DataValue::Null => return Ok(DataValue::Null),
98            _ => {
99                return Err(anyhow::anyhow!(
100                    "LPAD length must be a number, got {:?}",
101                    args[1]
102                ))
103            }
104        };
105
106        let pad_str = if args.len() == 3 {
107            match &args[2] {
108                DataValue::String(s) => {
109                    if s.is_empty() {
110                        return Err(anyhow::anyhow!("LPAD pad string cannot be empty"));
111                    }
112                    s.clone()
113                }
114                DataValue::Null => return Ok(DataValue::Null),
115                _ => args[2].to_string(),
116            }
117        } else {
118            " ".to_string()
119        };
120
121        if string.len() >= length {
122            // Truncate if string is longer than target length
123            Ok(DataValue::String(string.chars().take(length).collect()))
124        } else {
125            let pad_needed = length - string.len();
126            let pad_chars: Vec<char> = pad_str.chars().collect();
127            let mut result = String::with_capacity(length);
128
129            // Add padding
130            let full_pads = pad_needed / pad_chars.len();
131            let partial_pad = pad_needed % pad_chars.len();
132
133            for _ in 0..full_pads {
134                result.push_str(&pad_str);
135            }
136            for i in 0..partial_pad {
137                result.push(pad_chars[i]);
138            }
139
140            // Add original string
141            result.push_str(&string);
142
143            Ok(DataValue::String(result))
144        }
145    }
146}
147
148/// RPAD function - Right pad a string to a certain length
149pub struct RPadFunction;
150
151impl SqlFunction for RPadFunction {
152    fn signature(&self) -> FunctionSignature {
153        FunctionSignature {
154            name: "RPAD",
155            category: FunctionCategory::String,
156            arg_count: ArgCount::Range(2, 3),
157            description: "Right pad a string to a certain length",
158            returns: "String padded on the right to specified length",
159            examples: vec![
160                "SELECT RPAD('5', 3, '0')  -- Returns '500'",
161                "SELECT RPAD('hello', 10, '.')  -- Returns 'hello.....'",
162                "SELECT RPAD('abc', 5)  -- Returns 'abc  ' (default pad is space)",
163            ],
164        }
165    }
166
167    fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
168        if args.len() < 2 || args.len() > 3 {
169            return Err(anyhow::anyhow!("RPAD expects 2 or 3 arguments"));
170        }
171
172        let string = match &args[0] {
173            DataValue::String(s) => s.clone(),
174            DataValue::Null => return Ok(DataValue::Null),
175            _ => args[0].to_string(),
176        };
177
178        let length = match &args[1] {
179            DataValue::Integer(n) => *n as usize,
180            DataValue::Float(f) => *f as usize,
181            DataValue::Null => return Ok(DataValue::Null),
182            _ => {
183                return Err(anyhow::anyhow!(
184                    "RPAD length must be a number, got {:?}",
185                    args[1]
186                ))
187            }
188        };
189
190        let pad_str = if args.len() == 3 {
191            match &args[2] {
192                DataValue::String(s) => {
193                    if s.is_empty() {
194                        return Err(anyhow::anyhow!("RPAD pad string cannot be empty"));
195                    }
196                    s.clone()
197                }
198                DataValue::Null => return Ok(DataValue::Null),
199                _ => args[2].to_string(),
200            }
201        } else {
202            " ".to_string()
203        };
204
205        if string.len() >= length {
206            // Truncate if string is longer than target length
207            Ok(DataValue::String(string.chars().take(length).collect()))
208        } else {
209            let pad_needed = length - string.len();
210            let pad_chars: Vec<char> = pad_str.chars().collect();
211            let mut result = String::with_capacity(length);
212
213            // Add original string
214            result.push_str(&string);
215
216            // Add padding
217            let full_pads = pad_needed / pad_chars.len();
218            let partial_pad = pad_needed % pad_chars.len();
219
220            for _ in 0..full_pads {
221                result.push_str(&pad_str);
222            }
223            for i in 0..partial_pad {
224                result.push(pad_chars[i]);
225            }
226
227            Ok(DataValue::String(result))
228        }
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn test_repeat() {
238        let func = RepeatFunction;
239
240        // Basic repetition
241        assert_eq!(
242            func.evaluate(&[DataValue::String("*".to_string()), DataValue::Integer(5)])
243                .unwrap(),
244            DataValue::String("*****".to_string())
245        );
246
247        // Multi-char repetition
248        assert_eq!(
249            func.evaluate(&[DataValue::String("ab".to_string()), DataValue::Integer(3)])
250                .unwrap(),
251            DataValue::String("ababab".to_string())
252        );
253
254        // Zero repetitions
255        assert_eq!(
256            func.evaluate(&[DataValue::String("x".to_string()), DataValue::Integer(0)])
257                .unwrap(),
258            DataValue::String("".to_string())
259        );
260
261        // Null handling
262        assert_eq!(
263            func.evaluate(&[DataValue::Null, DataValue::Integer(5)])
264                .unwrap(),
265            DataValue::Null
266        );
267    }
268
269    #[test]
270    fn test_lpad() {
271        let func = LPadFunction;
272
273        // Basic padding with zeros
274        assert_eq!(
275            func.evaluate(&[
276                DataValue::String("5".to_string()),
277                DataValue::Integer(3),
278                DataValue::String("0".to_string())
279            ])
280            .unwrap(),
281            DataValue::String("005".to_string())
282        );
283
284        // Default space padding
285        assert_eq!(
286            func.evaluate(&[DataValue::String("abc".to_string()), DataValue::Integer(5)])
287                .unwrap(),
288            DataValue::String("  abc".to_string())
289        );
290
291        // Multi-char padding
292        assert_eq!(
293            func.evaluate(&[
294                DataValue::String("X".to_string()),
295                DataValue::Integer(5),
296                DataValue::String("ab".to_string())
297            ])
298            .unwrap(),
299            DataValue::String("ababX".to_string())
300        );
301    }
302
303    #[test]
304    fn test_rpad() {
305        let func = RPadFunction;
306
307        // Basic padding with zeros
308        assert_eq!(
309            func.evaluate(&[
310                DataValue::String("5".to_string()),
311                DataValue::Integer(3),
312                DataValue::String("0".to_string())
313            ])
314            .unwrap(),
315            DataValue::String("500".to_string())
316        );
317
318        // Default space padding
319        assert_eq!(
320            func.evaluate(&[DataValue::String("abc".to_string()), DataValue::Integer(5)])
321                .unwrap(),
322            DataValue::String("abc  ".to_string())
323        );
324
325        // Dots padding
326        assert_eq!(
327            func.evaluate(&[
328                DataValue::String("hello".to_string()),
329                DataValue::Integer(10),
330                DataValue::String(".".to_string())
331            ])
332            .unwrap(),
333            DataValue::String("hello.....".to_string())
334        );
335    }
336}
sql_cli/sql/functions/string_utils.rs

sql_cli/sql/functions/
string_utils.rs
