shape_runtime/intrinsics/
array_transforms.rs

1//! Column intrinsic functions
2
3use super::{
4    extract_f64, extract_f64_array, extract_usize, f64_vec_to_nb_array, i64_vec_to_nb_int_array,
5    option_i64_vec_to_nb, try_extract_i64_slice,
6};
7use crate::context::ExecutionContext;
8use crate::simd_i64;
9use shape_ast::error::{Result, ShapeError};
10use shape_value::ValueWord;
11
12pub fn intrinsic_column_select(
13    args: &[ValueWord],
14    _ctx: &mut ExecutionContext,
15) -> Result<ValueWord> {
16    if args.len() != 1 {
17        return Err(ShapeError::RuntimeError {
18            message: "__intrinsic_series requires exactly 1 argument".to_string(),
19            location: None,
20        });
21    }
22
23    let data = extract_f64_array(&args[0], "series")?;
24    Ok(f64_vec_to_nb_array(data))
25}
26
27pub fn intrinsic_shift(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
28    if args.len() != 2 {
29        return Err(ShapeError::RuntimeError {
30            message: "shift() requires 2 arguments (series, shift)".to_string(),
31            location: None,
32        });
33    }
34
35    let data = extract_f64_array(&args[0], "shift")?;
36    let shift = extract_f64(&args[1], "shift() second argument")? as i64;
37
38    let mut result = Vec::with_capacity(data.len());
39
40    if shift > 0 {
41        let shift_usize = shift.min(data.len() as i64) as usize;
42        for _ in 0..shift_usize {
43            result.push(f64::NAN);
44        }
45        result.extend_from_slice(&data[..data.len().saturating_sub(shift_usize)]);
46    } else if shift < 0 {
47        let shift_usize = (-shift).min(data.len() as i64) as usize;
48        result.extend_from_slice(&data[shift_usize..]);
49        for _ in 0..shift_usize {
50            result.push(f64::NAN);
51        }
52    } else {
53        result.extend_from_slice(&data);
54    }
55
56    Ok(f64_vec_to_nb_array(result))
57}
58
59pub fn intrinsic_diff(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
60    if args.is_empty() || args.len() > 2 {
61        return Err(ShapeError::RuntimeError {
62            message: "diff() requires 1 or 2 arguments (series, [period])".to_string(),
63            location: None,
64        });
65    }
66
67    let period = if args.len() == 2 {
68        extract_usize(&args[1], "diff() second argument")?
69    } else {
70        1
71    };
72
73    // i64 fast path
74    if let Some(slice) = try_extract_i64_slice(&args[0]) {
75        return Ok(option_i64_vec_to_nb(simd_i64::diff_i64(slice, period)));
76    }
77
78    let data = extract_f64_array(&args[0], "diff")?;
79
80    let mut result = Vec::with_capacity(data.len());
81    for i in 0..data.len() {
82        if i < period {
83            result.push(f64::NAN);
84        } else {
85            result.push(data[i] - data[i - period]);
86        }
87    }
88
89    Ok(f64_vec_to_nb_array(result))
90}
91
92pub fn intrinsic_pct_change(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
93    if args.is_empty() || args.len() > 2 {
94        return Err(ShapeError::RuntimeError {
95            message: "pct_change() requires 1 or 2 arguments (series, [period])".to_string(),
96            location: None,
97        });
98    }
99
100    let data = extract_f64_array(&args[0], "pct_change")?;
101
102    let period = if args.len() == 2 {
103        extract_usize(&args[1], "pct_change() second argument")?
104    } else {
105        1
106    };
107
108    let mut result = Vec::with_capacity(data.len());
109    for i in 0..data.len() {
110        if i < period {
111            result.push(f64::NAN);
112        } else {
113            let prev = data[i - period];
114            if prev == 0.0 {
115                result.push(f64::NAN);
116            } else {
117                result.push((data[i] - prev) / prev);
118            }
119        }
120    }
121
122    Ok(f64_vec_to_nb_array(result))
123}
124
125pub fn intrinsic_fillna(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
126    if args.len() != 2 {
127        return Err(ShapeError::RuntimeError {
128            message: "fillna() requires 2 arguments (series, value)".to_string(),
129            location: None,
130        });
131    }
132
133    let data = extract_f64_array(&args[0], "fillna")?;
134    let fill_value = extract_f64(&args[1], "fillna() second argument")?;
135
136    let result: Vec<f64> = data
137        .iter()
138        .map(|&v| if v.is_nan() { fill_value } else { v })
139        .collect();
140
141    Ok(f64_vec_to_nb_array(result))
142}
143
144pub fn intrinsic_cumsum(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
145    if args.len() != 1 {
146        return Err(ShapeError::RuntimeError {
147            message: "cumsum() requires 1 argument (series)".to_string(),
148            location: None,
149        });
150    }
151
152    // i64 fast path
153    if let Some(slice) = try_extract_i64_slice(&args[0]) {
154        return Ok(i64_vec_to_nb_int_array(simd_i64::cumsum_i64(slice)));
155    }
156
157    let data = extract_f64_array(&args[0], "cumsum")?;
158    let mut result = Vec::with_capacity(data.len());
159    let mut acc = 0.0;
160    for &v in &data {
161        acc += v;
162        result.push(acc);
163    }
164
165    Ok(f64_vec_to_nb_array(result))
166}
167
168pub fn intrinsic_cumprod(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
169    if args.len() != 1 {
170        return Err(ShapeError::RuntimeError {
171            message: "cumprod() requires 1 argument (series)".to_string(),
172            location: None,
173        });
174    }
175
176    let data = extract_f64_array(&args[0], "cumprod")?;
177    let mut result = Vec::with_capacity(data.len());
178    let mut acc = 1.0;
179    for &v in &data {
180        acc *= v;
181        result.push(acc);
182    }
183
184    Ok(f64_vec_to_nb_array(result))
185}
186
187pub fn intrinsic_clip(args: &[ValueWord], _ctx: &mut ExecutionContext) -> Result<ValueWord> {
188    if args.len() != 3 {
189        return Err(ShapeError::RuntimeError {
190            message: "clip() requires 3 arguments (series, min, max)".to_string(),
191            location: None,
192        });
193    }
194
195    let data = extract_f64_array(&args[0], "clip")?;
196    let min = extract_f64(&args[1], "clip() second argument")?;
197    let max = extract_f64(&args[2], "clip() third argument")?;
198
199    let result: Vec<f64> = data.iter().map(|&v| v.max(min).min(max)).collect();
200
201    Ok(f64_vec_to_nb_array(result))
202}
shape_runtime/intrinsics/array_transforms.rs

shape_runtime/intrinsics/
array_transforms.rs
