alopex_core/dataframe/
mod.rs

1//! DataFrame-oriented primitives shared by higher-level crates.
2
3use regex::Regex;
4
5use crate::{Error, Result};
6
7pub mod cast;
8pub mod partition_scan;
9
10const MICROS_PER_SECOND: i64 = 1_000_000;
11const SECONDS_PER_DAY: i64 = 86_400;
12
13/// Nullable string column used by core DataFrame namespace primitives.
14pub type Utf8Column = Vec<Option<String>>;
15/// Nullable boolean column used by core DataFrame namespace primitives.
16pub type BoolColumn = Vec<Option<bool>>;
17/// Nullable unsigned integer column used by core DataFrame namespace primitives.
18pub type UIntColumn = Vec<Option<usize>>;
19/// Nullable timestamp column using microseconds since Unix epoch.
20pub type TimestampMicrosColumn = Vec<Option<i64>>;
21/// Nullable list column with nullable elements.
22pub type ListColumn<T> = Vec<Option<Vec<Option<T>>>>;
23
24/// Result of exploding one list column.
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct ExplodedColumn<T> {
27    /// Exploded values. Null and empty source lists both produce one null row.
28    pub values: Vec<Option<T>>,
29    /// Source row index for each exploded value.
30    pub source_rows: Vec<usize>,
31}
32
33/// Convert UTF-8 strings to lowercase while preserving nulls.
34pub fn str_to_lowercase(values: &[Option<String>]) -> Utf8Column {
35    values
36        .iter()
37        .map(|value| value.as_ref().map(|text| text.to_lowercase()))
38        .collect()
39}
40
41/// Convert UTF-8 strings to uppercase while preserving nulls.
42pub fn str_to_uppercase(values: &[Option<String>]) -> Utf8Column {
43    values
44        .iter()
45        .map(|value| value.as_ref().map(|text| text.to_uppercase()))
46        .collect()
47}
48
49/// Return whether each UTF-8 value matches a regular expression pattern.
50pub fn str_contains(values: &[Option<String>], pattern: &str) -> Result<BoolColumn> {
51    let regex = compile_regex("str.contains", pattern)?;
52    Ok(values
53        .iter()
54        .map(|value| value.as_ref().map(|text| regex.is_match(text)))
55        .collect())
56}
57
58/// Replace all regular expression matches in each UTF-8 value.
59pub fn str_replace(
60    values: &[Option<String>],
61    pattern: &str,
62    replacement: &str,
63) -> Result<Utf8Column> {
64    let regex = compile_regex("str.replace", pattern)?;
65    Ok(values
66        .iter()
67        .map(|value| {
68            value
69                .as_ref()
70                .map(|text| regex.replace_all(text, replacement).into_owned())
71        })
72        .collect())
73}
74
75/// Strip characters from both ends of each UTF-8 value.
76pub fn str_strip_chars(values: &[Option<String>], chars: Option<&str>) -> Utf8Column {
77    values
78        .iter()
79        .map(|value| {
80            value.as_ref().map(|text| match chars {
81                Some(chars) => text.trim_matches(|ch| chars.contains(ch)).to_string(),
82                None => text.trim().to_string(),
83            })
84        })
85        .collect()
86}
87
88/// Split each UTF-8 value by a literal separator.
89pub fn str_split(values: &[Option<String>], separator: &str) -> ListColumn<String> {
90    values
91        .iter()
92        .map(|value| {
93            value.as_ref().map(|text| {
94                text.split(separator)
95                    .map(|part| Some(part.to_string()))
96                    .collect()
97            })
98        })
99        .collect()
100}
101
102/// Count Unicode scalar values in each UTF-8 value.
103pub fn str_len_chars(values: &[Option<String>]) -> UIntColumn {
104    values
105        .iter()
106        .map(|value| value.as_ref().map(|text| text.chars().count()))
107        .collect()
108}
109
110/// Extract the first regular expression match or capture group from each value.
111pub fn str_extract(
112    values: &[Option<String>],
113    pattern: &str,
114    capture_group: usize,
115) -> Result<Utf8Column> {
116    let regex = compile_regex("str.extract", pattern)?;
117    Ok(values
118        .iter()
119        .map(|value| {
120            value.as_ref().and_then(|text| {
121                regex.captures(text).and_then(|captures| {
122                    captures
123                        .get(capture_group)
124                        .map(|matched| matched.as_str().to_string())
125                })
126            })
127        })
128        .collect())
129}
130
131/// Extract UTC year from timestamp micros.
132pub fn dt_year(values: &[Option<i64>]) -> Vec<Option<i32>> {
133    values
134        .iter()
135        .map(|value| value.map(timestamp_parts).map(|parts| parts.year))
136        .collect()
137}
138
139/// Extract UTC month from timestamp micros, in the range 1..=12.
140pub fn dt_month(values: &[Option<i64>]) -> Vec<Option<u32>> {
141    values
142        .iter()
143        .map(|value| value.map(timestamp_parts).map(|parts| parts.month))
144        .collect()
145}
146
147/// Extract UTC day of month from timestamp micros, in the range 1..=31.
148pub fn dt_day(values: &[Option<i64>]) -> Vec<Option<u32>> {
149    values
150        .iter()
151        .map(|value| value.map(timestamp_parts).map(|parts| parts.day))
152        .collect()
153}
154
155/// Extract ISO weekday from timestamp micros, where Monday is 1 and Sunday is 7.
156pub fn dt_weekday(values: &[Option<i64>]) -> Vec<Option<u32>> {
157    values
158        .iter()
159        .map(|value| value.map(|micros| iso_weekday(timestamp_days(micros))))
160        .collect()
161}
162
163/// Format timestamp micros as UTC RFC3339-like text.
164pub fn dt_to_string(values: &[Option<i64>]) -> Utf8Column {
165    values
166        .iter()
167        .map(|value| value.map(format_timestamp_micros))
168        .collect()
169}
170
171/// Convert naive local timestamp micros between fixed-offset time zones.
172///
173/// Offsets use `Z`, `+HH:MM`, or `-HH:MM`. The returned value is the target
174/// local timestamp represented with the same microsecond epoch encoding.
175pub fn dt_convert_time_zone(
176    values: &[Option<i64>],
177    from_offset: &str,
178    to_offset: &str,
179) -> Result<TimestampMicrosColumn> {
180    let from = parse_offset_seconds(from_offset)?;
181    let to = parse_offset_seconds(to_offset)?;
182    let delta_micros = i64::from(to - from)
183        .checked_mul(MICROS_PER_SECOND)
184        .ok_or_else(|| invalid_parameter("dt.convert_time_zone", "offset delta overflow"))?;
185    values
186        .iter()
187        .map(|value| {
188            value
189                .map(|micros| {
190                    micros.checked_add(delta_micros).ok_or_else(|| {
191                        invalid_parameter("dt.convert_time_zone", "timestamp overflow")
192                    })
193                })
194                .transpose()
195        })
196        .collect()
197}
198
199/// Join string list elements with a separator.
200///
201/// Null lists return null. Null elements use `null_value` when provided;
202/// otherwise the whole row returns null.
203pub fn list_join(
204    values: &[Option<Vec<Option<String>>>],
205    separator: &str,
206    null_value: Option<&str>,
207) -> Utf8Column {
208    values
209        .iter()
210        .map(|list| {
211            list.as_ref().and_then(|items| {
212                let mut parts = Vec::with_capacity(items.len());
213                for item in items {
214                    match item {
215                        Some(text) => parts.push(text.as_str()),
216                        None => parts.push(null_value?),
217                    }
218                }
219                Some(parts.join(separator))
220            })
221        })
222        .collect()
223}
224
225/// Return list lengths while preserving null lists.
226pub fn list_len<T>(values: &[Option<Vec<Option<T>>>]) -> UIntColumn {
227    values
228        .iter()
229        .map(|list| list.as_ref().map(Vec::len))
230        .collect()
231}
232
233/// Return whether each list contains a non-null needle value.
234pub fn list_contains<T: PartialEq>(values: &[Option<Vec<Option<T>>>], needle: &T) -> BoolColumn {
235    values
236        .iter()
237        .map(|list| {
238            list.as_ref()
239                .map(|items| items.iter().any(|item| item.as_ref() == Some(needle)))
240        })
241        .collect()
242}
243
244/// Explode one list column into values and source row indexes.
245pub fn explode_list<T: Clone>(values: &[Option<Vec<Option<T>>>]) -> ExplodedColumn<T> {
246    let mut exploded = Vec::new();
247    let mut source_rows = Vec::new();
248    for (row_idx, list) in values.iter().enumerate() {
249        match list {
250            Some(items) if !items.is_empty() => {
251                for item in items {
252                    exploded.push(item.clone());
253                    source_rows.push(row_idx);
254                }
255            }
256            Some(_) | None => {
257                exploded.push(None);
258                source_rows.push(row_idx);
259            }
260        }
261    }
262    ExplodedColumn {
263        values: exploded,
264        source_rows,
265    }
266}
267
268/// Implode values into nullable lists according to contiguous group lengths.
269pub fn implode_by_group_lengths<T: Clone>(
270    values: &[Option<T>],
271    group_lengths: &[usize],
272) -> Result<ListColumn<T>> {
273    let expected = group_lengths
274        .iter()
275        .try_fold(0usize, |acc, len| acc.checked_add(*len))
276        .ok_or_else(|| invalid_parameter("df.implode", "group length overflow"))?;
277    if expected != values.len() {
278        return Err(invalid_parameter(
279            "df.implode",
280            format!("group lengths sum to {expected}, expected {}", values.len()),
281        ));
282    }
283
284    let mut offset = 0usize;
285    let mut output = Vec::with_capacity(group_lengths.len());
286    for &len in group_lengths {
287        let end = offset + len;
288        output.push(Some(values[offset..end].to_vec()));
289        offset = end;
290    }
291    Ok(output)
292}
293
294fn compile_regex(operation: &str, pattern: &str) -> Result<Regex> {
295    Regex::new(pattern).map_err(|err| invalid_parameter(operation, err.to_string()))
296}
297
298fn invalid_parameter(param: impl Into<String>, reason: impl Into<String>) -> Error {
299    Error::InvalidParameter {
300        param: param.into(),
301        reason: reason.into(),
302    }
303}
304
305#[derive(Clone, Copy)]
306struct TimestampParts {
307    year: i32,
308    month: u32,
309    day: u32,
310    hour: i64,
311    minute: i64,
312    second: i64,
313    micros: i64,
314}
315
316fn timestamp_parts(micros: i64) -> TimestampParts {
317    let seconds = micros.div_euclid(MICROS_PER_SECOND);
318    let micros_remainder = micros.rem_euclid(MICROS_PER_SECOND);
319    let days = seconds.div_euclid(SECONDS_PER_DAY);
320    let second_of_day = seconds.rem_euclid(SECONDS_PER_DAY);
321    let (year, month, day) = civil_from_days(days);
322    TimestampParts {
323        year,
324        month,
325        day,
326        hour: second_of_day / 3_600,
327        minute: (second_of_day % 3_600) / 60,
328        second: second_of_day % 60,
329        micros: micros_remainder,
330    }
331}
332
333fn timestamp_days(micros: i64) -> i64 {
334    micros
335        .div_euclid(MICROS_PER_SECOND)
336        .div_euclid(SECONDS_PER_DAY)
337}
338
339fn iso_weekday(days_since_epoch: i64) -> u32 {
340    (days_since_epoch + 3).rem_euclid(7) as u32 + 1
341}
342
343fn format_timestamp_micros(micros: i64) -> String {
344    let parts = timestamp_parts(micros);
345    if parts.micros == 0 {
346        format!(
347            "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
348            parts.year, parts.month, parts.day, parts.hour, parts.minute, parts.second
349        )
350    } else {
351        let mut fraction = format!("{:06}", parts.micros);
352        while fraction.ends_with('0') {
353            fraction.pop();
354        }
355        format!(
356            "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{}Z",
357            parts.year, parts.month, parts.day, parts.hour, parts.minute, parts.second, fraction
358        )
359    }
360}
361
362fn parse_offset_seconds(offset: &str) -> Result<i32> {
363    if offset == "Z" || offset == "+00:00" || offset == "-00:00" {
364        return Ok(0);
365    }
366    let bytes = offset.as_bytes();
367    if bytes.len() != 6 || bytes[3] != b':' || !matches!(bytes[0], b'+' | b'-') {
368        return Err(invalid_parameter(
369            "dt.convert_time_zone",
370            "offset must be Z, +HH:MM, or -HH:MM",
371        ));
372    }
373    let hours = offset[1..3].parse::<i32>().map_err(|_| {
374        invalid_parameter(
375            "dt.convert_time_zone",
376            "offset hour must be two decimal digits",
377        )
378    })?;
379    let minutes = offset[4..6].parse::<i32>().map_err(|_| {
380        invalid_parameter(
381            "dt.convert_time_zone",
382            "offset minute must be two decimal digits",
383        )
384    })?;
385    if hours > 23 || minutes > 59 {
386        return Err(invalid_parameter(
387            "dt.convert_time_zone",
388            "offset is out of range",
389        ));
390    }
391    let seconds = hours * 3_600 + minutes * 60;
392    Ok(if bytes[0] == b'-' { -seconds } else { seconds })
393}
394
395fn civil_from_days(days: i64) -> (i32, u32, u32) {
396    let days = days + 719_468;
397    let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
398    let day_of_era = days - era * 146_097;
399    let year_of_era =
400        (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
401    let mut year = year_of_era + era * 400;
402    let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
403    let month_prime = (5 * day_of_year + 2) / 153;
404    let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
405    let month = month_prime + if month_prime < 10 { 3 } else { -9 };
406    year += i64::from(month <= 2);
407    (year as i32, month as u32, day as u32)
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    fn strings(values: &[Option<&str>]) -> Vec<Option<String>> {
415        values
416            .iter()
417            .map(|value| value.map(str::to_string))
418            .collect()
419    }
420
421    #[test]
422    fn string_namespace_primitives_preserve_nulls_and_unicode() {
423        let input = strings(&[Some(" Alopex "), None, Some("Straße42")]);
424
425        assert_eq!(
426            str_to_lowercase(&input),
427            strings(&[Some(" alopex "), None, Some("straße42")])
428        );
429        assert_eq!(
430            str_to_uppercase(&input),
431            strings(&[Some(" ALOPEX "), None, Some("STRASSE42")])
432        );
433        assert_eq!(
434            str_strip_chars(&input, None),
435            strings(&[Some("Alopex"), None, Some("Straße42")])
436        );
437        assert_eq!(str_len_chars(&input), vec![Some(8), None, Some(8)]);
438        assert_eq!(
439            str_contains(&input, r"\d+$").unwrap(),
440            vec![Some(false), None, Some(true)]
441        );
442        assert_eq!(
443            str_replace(&input, r"\d+", "#").unwrap(),
444            strings(&[Some(" Alopex "), None, Some("Straße#")])
445        );
446        assert_eq!(
447            str_extract(&input, r"(\p{Alphabetic}+)(\d+)", 1).unwrap(),
448            strings(&[None, None, Some("Straße")])
449        );
450    }
451
452    #[test]
453    fn string_split_returns_nullable_list_column() {
454        let input = strings(&[Some("a,b,"), None, Some("one")]);
455
456        assert_eq!(
457            str_split(&input, ","),
458            vec![
459                Some(strings(&[Some("a"), Some("b"), Some("")])),
460                None,
461                Some(strings(&[Some("one")]))
462            ]
463        );
464    }
465
466    #[test]
467    fn datetime_primitives_extract_and_format_utc_parts() {
468        let input = vec![Some(0), Some(1_704_067_200_123_000), None];
469
470        assert_eq!(dt_year(&input), vec![Some(1970), Some(2024), None]);
471        assert_eq!(dt_month(&input), vec![Some(1), Some(1), None]);
472        assert_eq!(dt_day(&input), vec![Some(1), Some(1), None]);
473        assert_eq!(dt_weekday(&input), vec![Some(4), Some(1), None]);
474        assert_eq!(
475            dt_to_string(&input),
476            strings(&[
477                Some("1970-01-01T00:00:00Z"),
478                Some("2024-01-01T00:00:00.123Z"),
479                None
480            ])
481        );
482    }
483
484    #[test]
485    fn timezone_conversion_uses_fixed_offsets_and_checks_errors() {
486        let input = vec![Some(0), None];
487
488        assert_eq!(
489            dt_convert_time_zone(&input, "Z", "+09:00").unwrap(),
490            vec![Some(32_400_000_000), None]
491        );
492        assert!(matches!(
493            dt_convert_time_zone(&input, "UTC", "+09:00"),
494            Err(Error::InvalidParameter { param, .. }) if param == "dt.convert_time_zone"
495        ));
496    }
497
498    #[test]
499    fn list_namespace_primitives_handle_null_lists_and_elements() {
500        let input = vec![
501            Some(vec![Some("a".to_string()), None, Some("c".to_string())]),
502            None,
503            Some(Vec::new()),
504        ];
505
506        assert_eq!(
507            list_join(&input, "-", Some("NULL")),
508            strings(&[Some("a-NULL-c"), None, Some("")])
509        );
510        assert_eq!(
511            list_join(&input, "-", None),
512            strings(&[None, None, Some("")])
513        );
514        assert_eq!(list_len(&input), vec![Some(3), None, Some(0)]);
515        assert_eq!(
516            list_contains(&input, &"c".to_string()),
517            vec![Some(true), None, Some(false)]
518        );
519    }
520
521    #[test]
522    fn explode_and_implode_primitives_are_deterministic() {
523        let input = vec![Some(vec![Some(1), None, Some(3)]), Some(Vec::new()), None];
524
525        let exploded = explode_list(&input);
526        assert_eq!(exploded.values, vec![Some(1), None, Some(3), None, None]);
527        assert_eq!(exploded.source_rows, vec![0, 0, 0, 1, 2]);
528        assert_eq!(
529            implode_by_group_lengths(&exploded.values, &[3, 1, 1]).unwrap(),
530            vec![
531                Some(vec![Some(1), None, Some(3)]),
532                Some(vec![None]),
533                Some(vec![None])
534            ]
535        );
536        assert!(matches!(
537            implode_by_group_lengths(&exploded.values, &[2]),
538            Err(Error::InvalidParameter { param, .. }) if param == "df.implode"
539        ));
540    }
541}
alopex_core/dataframe/mod.rs

alopex_core/dataframe/
mod.rs
