linkrs_query/executor/expression/evaluator/
functions.rs1use crate::executor::expression::{ExpressionError, ExpressionErrorType};
6use linkrs_core::types::operators::AggregateFunction;
7use linkrs_core::value::list::List;
8use linkrs_core::value::NullType;
9use linkrs_core::Value;
10
11pub struct FunctionEvaluator;
13
14impl FunctionEvaluator {
15 pub fn eval_aggregate_function(
17 func: &AggregateFunction,
18 args: &[Value],
19 distinct: bool,
20 ) -> Result<Value, ExpressionError> {
21 if args.is_empty() {
22 return Err(ExpressionError::argument_count_error(1, 0));
23 }
24
25 match func {
26 AggregateFunction::Count => {
27 if distinct {
28 let unique_values: std::collections::HashSet<_> = args.iter().collect();
29 Ok(Value::BigInt(unique_values.len() as i64))
30 } else {
31 Ok(Value::BigInt(args.len() as i64))
32 }
33 }
34 AggregateFunction::Sum => {
35 let mut sum = Value::Int(0);
36 for arg in args {
37 sum = sum.add(arg).map_err(ExpressionError::runtime_error)?;
38 }
39 Ok(sum)
40 }
41 AggregateFunction::Avg => {
42 let sum = Self::eval_aggregate_function(&AggregateFunction::Sum, args, distinct)?;
43 let count =
44 Self::eval_aggregate_function(&AggregateFunction::Count, args, distinct)?;
45 sum.div(&count).map_err(ExpressionError::runtime_error)
46 }
47 AggregateFunction::Min => {
48 let mut min = args[0].clone();
49 for arg in args.iter().skip(1) {
50 if arg < &min {
51 min = arg.clone();
52 }
53 }
54 Ok(min)
55 }
56 AggregateFunction::Max => {
57 let mut max = args[0].clone();
58 for arg in args.iter().skip(1) {
59 if arg > &max {
60 max = arg.clone();
61 }
62 }
63 Ok(max)
64 }
65 AggregateFunction::Collect => {
66 if distinct {
67 let unique_values: std::collections::HashSet<_> =
68 args.iter().cloned().collect();
69 Ok(Value::list(List::from(
70 unique_values.into_iter().collect::<Vec<_>>(),
71 )))
72 } else {
73 Ok(Value::list(List::from(args.to_vec())))
74 }
75 }
76 AggregateFunction::CollectSet => {
77 let unique_values: std::collections::HashSet<_> = args.iter().cloned().collect();
78 Ok(Value::set(unique_values))
79 }
80 AggregateFunction::Percentile => {
81 if args.len() < 2 {
82 return Err(ExpressionError::argument_count_error(2, args.len()));
83 }
84
85 let percentile = match &args[1] {
86 Value::Int(v) => *v as f64,
87 Value::BigInt(v) => *v as f64,
88 Value::Float(v) => *v as f64,
89 Value::Double(v) => *v,
90 _ => return Err(ExpressionError::type_error("Percentile must be a number")),
91 };
92
93 if !(0.0..=100.0).contains(&percentile) {
94 return Err(ExpressionError::new(
95 ExpressionErrorType::InvalidOperation,
96 "Percentile must be between 0 and 100",
97 ));
98 }
99
100 let values = match &args[0] {
101 Value::List(list) => list,
102 _ => return Err(ExpressionError::type_error("First argument must be a list")),
103 };
104
105 if values.is_empty() {
106 return Ok(Value::Null(linkrs_core::NullType::NaN));
107 }
108
109 let mut numeric_values = Vec::new();
110 for value in values.iter() {
111 match value {
112 Value::SmallInt(v) => numeric_values.push(*v as f64),
113 Value::Int(v) => numeric_values.push(*v as f64),
114 Value::BigInt(v) => numeric_values.push(*v as f64),
115 Value::Float(v) => numeric_values.push(*v as f64),
116 Value::Double(v) => numeric_values.push(*v),
117 _ => continue,
118 }
119 }
120
121 if numeric_values.is_empty() {
122 return Ok(Value::Null(linkrs_core::NullType::NaN));
123 }
124
125 numeric_values
126 .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
127
128 let index = (percentile / 100.0) * (numeric_values.len() - 1) as f64;
129 let lower_index = index.floor() as usize;
130 let upper_index = index.ceil() as usize;
131
132 if lower_index == upper_index {
133 Ok(Value::Double(numeric_values[lower_index]))
134 } else {
135 let lower_value = numeric_values[lower_index];
136 let upper_value = numeric_values[upper_index];
137 let weight = index - lower_index as f64;
138 let interpolated = lower_value + weight * (upper_value - lower_value);
139 Ok(Value::Double(interpolated))
140 }
141 }
142 AggregateFunction::Std => {
143 if args.is_empty() {
144 return Err(ExpressionError::argument_count_error(1, args.len()));
145 }
146
147 let values = match &args[0] {
148 Value::List(list) => list,
149 _ => return Err(ExpressionError::type_error("First argument must be a list")),
150 };
151
152 if values.is_empty() {
153 return Ok(Value::Null(linkrs_core::NullType::NaN));
154 }
155
156 let mut numeric_values = Vec::new();
157 for value in values.iter() {
158 match value {
159 Value::SmallInt(v) => numeric_values.push(*v as f64),
160 Value::Int(v) => numeric_values.push(*v as f64),
161 Value::BigInt(v) => numeric_values.push(*v as f64),
162 Value::Float(v) => numeric_values.push(*v as f64),
163 Value::Double(v) => numeric_values.push(*v),
164 _ => continue,
165 }
166 }
167
168 if numeric_values.is_empty() {
169 return Ok(Value::Null(linkrs_core::NullType::NaN));
170 }
171
172 let n = numeric_values.len() as f64;
173 let mean: f64 = numeric_values.iter().sum::<f64>() / n;
174 let variance: f64 = numeric_values
175 .iter()
176 .map(|value| (value - mean).powi(2))
177 .sum::<f64>()
178 / n;
179 let std_dev = variance.sqrt();
180
181 Ok(Value::Double(std_dev))
182 }
183 AggregateFunction::BitAnd => {
184 if args.is_empty() {
185 return Err(ExpressionError::argument_count_error(1, args.len()));
186 }
187
188 let values = match &args[0] {
189 Value::List(list) => list,
190 _ => return Err(ExpressionError::type_error("First argument must be a list")),
191 };
192
193 if values.is_empty() {
194 return Ok(Value::Null(linkrs_core::NullType::NaN));
195 }
196
197 let mut result = i64::MAX;
198 for value in values.iter() {
199 match value {
200 Value::SmallInt(v) => result &= *v as i64,
201 Value::Int(v) => result &= *v as i64,
202 Value::BigInt(v) => result &= *v,
203 _ => {
204 return Err(ExpressionError::type_error(
205 "All values must be integers for BIT_AND",
206 ))
207 }
208 }
209 }
210
211 Ok(Value::BigInt(result))
212 }
213 AggregateFunction::BitOr => {
214 if args.is_empty() {
215 return Err(ExpressionError::argument_count_error(1, args.len()));
216 }
217
218 let values = match &args[0] {
219 Value::List(list) => list,
220 _ => return Err(ExpressionError::type_error("First argument must be a list")),
221 };
222
223 if values.is_empty() {
224 return Ok(Value::Null(linkrs_core::NullType::NaN));
225 }
226
227 let mut result = 0i64;
228 for value in values.iter() {
229 match value {
230 Value::SmallInt(v) => result |= *v as i64,
231 Value::Int(v) => result |= *v as i64,
232 Value::BigInt(v) => result |= *v,
233 _ => {
234 return Err(ExpressionError::type_error(
235 "All values must be integers for BIT_OR",
236 ))
237 }
238 }
239 }
240
241 Ok(Value::BigInt(result))
242 }
243 AggregateFunction::Variance => {
244 if args.is_empty() {
245 return Err(ExpressionError::argument_count_error(1, args.len()));
246 }
247
248 let values = match &args[0] {
249 Value::List(list) => list,
250 _ => return Err(ExpressionError::type_error("First argument must be a list")),
251 };
252
253 if values.is_empty() {
254 return Ok(Value::Null(linkrs_core::NullType::NaN));
255 }
256
257 let mut numeric_values = Vec::new();
258 for value in values.iter() {
259 match value {
260 Value::SmallInt(v) => numeric_values.push(*v as f64),
261 Value::Int(v) => numeric_values.push(*v as f64),
262 Value::BigInt(v) => numeric_values.push(*v as f64),
263 Value::Float(v) => numeric_values.push(*v as f64),
264 Value::Double(v) => numeric_values.push(*v),
265 _ => continue,
266 }
267 }
268
269 if numeric_values.is_empty() {
270 return Ok(Value::Null(linkrs_core::NullType::NaN));
271 }
272
273 let n = numeric_values.len() as f64;
274 let mean: f64 = numeric_values.iter().sum::<f64>() / n;
275 let variance: f64 = numeric_values
276 .iter()
277 .map(|value| (value - mean).powi(2))
278 .sum::<f64>()
279 / n;
280
281 Ok(Value::Double(variance))
282 }
283 AggregateFunction::Median => {
284 if args.is_empty() {
285 return Err(ExpressionError::argument_count_error(1, args.len()));
286 }
287
288 let values = match &args[0] {
289 Value::List(list) => list,
290 _ => return Err(ExpressionError::type_error("First argument must be a list")),
291 };
292
293 if values.is_empty() {
294 return Ok(Value::Null(linkrs_core::NullType::NaN));
295 }
296
297 let mut numeric_values = Vec::new();
298 for value in values.iter() {
299 match value {
300 Value::SmallInt(v) => numeric_values.push(*v as f64),
301 Value::Int(v) => numeric_values.push(*v as f64),
302 Value::BigInt(v) => numeric_values.push(*v as f64),
303 Value::Float(v) => numeric_values.push(*v as f64),
304 Value::Double(v) => numeric_values.push(*v),
305 _ => continue,
306 }
307 }
308
309 if numeric_values.is_empty() {
310 return Ok(Value::Null(linkrs_core::NullType::NaN));
311 }
312
313 numeric_values
314 .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
315
316 let len = numeric_values.len();
317 let median = if len % 2 == 0 {
318 (numeric_values[len / 2 - 1] + numeric_values[len / 2]) / 2.0
319 } else {
320 numeric_values[len / 2]
321 };
322
323 Ok(Value::Double(median))
324 }
325 AggregateFunction::Mode => {
326 if args.is_empty() {
327 return Err(ExpressionError::argument_count_error(1, args.len()));
328 }
329
330 let values = match &args[0] {
331 Value::List(list) => list,
332 _ => return Err(ExpressionError::type_error("First argument must be a list")),
333 };
334
335 if values.is_empty() {
336 return Ok(Value::Null(linkrs_core::NullType::NaN));
337 }
338
339 let mut frequency_map = std::collections::HashMap::new();
340 for value in values.iter() {
341 let key = format!("{}", value);
342 *frequency_map.entry(key).or_insert(0usize) += 1;
343 }
344
345 let mode = frequency_map
346 .into_iter()
347 .max_by_key(|(_, count)| *count)
348 .map(|(key, _)| key);
349
350 match mode {
351 Some(mode_str) => {
352 if let Ok(int_val) = mode_str.parse::<i32>() {
353 Ok(Value::Int(int_val))
354 } else if let Ok(float_val) = mode_str.parse::<f64>() {
355 Ok(Value::Double(float_val))
356 } else if mode_str == "true" {
357 Ok(Value::Bool(true))
358 } else if mode_str == "false" {
359 Ok(Value::Bool(false))
360 } else {
361 Ok(Value::string(mode_str))
362 }
363 }
364 None => Ok(Value::Null(linkrs_core::NullType::NaN)),
365 }
366 }
367 AggregateFunction::BoolAnd => {
368 if args.is_empty() {
369 return Err(ExpressionError::argument_count_error(1, args.len()));
370 }
371
372 let values = match &args[0] {
373 Value::List(list) => list,
374 _ => return Err(ExpressionError::type_error("First argument must be a list")),
375 };
376
377 if values.is_empty() {
378 return Ok(Value::Bool(true));
379 }
380
381 let mut result = true;
382 for value in values.iter() {
383 match value {
384 Value::Bool(b) => result &= *b,
385 Value::Null(_) => {}
386 _ => {
387 return Err(ExpressionError::type_error(
388 "All values must be boolean for BOOL_AND",
389 ))
390 }
391 }
392 }
393
394 Ok(Value::Bool(result))
395 }
396 AggregateFunction::BoolOr => {
397 if args.is_empty() {
398 return Err(ExpressionError::argument_count_error(1, args.len()));
399 }
400
401 let values = match &args[0] {
402 Value::List(list) => list,
403 _ => return Err(ExpressionError::type_error("First argument must be a list")),
404 };
405
406 if values.is_empty() {
407 return Ok(Value::Bool(false));
408 }
409
410 let mut result = false;
411 for value in values.iter() {
412 match value {
413 Value::Bool(b) => result |= *b,
414 Value::Null(_) => {}
415 _ => {
416 return Err(ExpressionError::type_error(
417 "All values must be boolean for BOOL_OR",
418 ))
419 }
420 }
421 }
422
423 Ok(Value::Bool(result))
424 }
425 AggregateFunction::GroupConcat => {
426 if args.is_empty() {
427 return Err(ExpressionError::argument_count_error(1, args.len()));
428 }
429
430 let values = match &args[0] {
431 Value::List(list) => list,
432 _ => return Err(ExpressionError::type_error("First argument must be a list")),
433 };
434
435 if values.is_empty() {
436 return Ok(Value::string(String::new()));
437 }
438
439 let separator = match args.get(1) {
440 Some(Value::String(s)) => s.to_string(),
441 _ => ",".to_string(),
442 };
443
444 let result: Vec<String> = values.iter().map(|v| format!("{}", v)).collect();
445 Ok(Value::string(result.join(&separator)))
446 }
447 AggregateFunction::GroupConcatWithOrder => {
448 if args.is_empty() {
449 return Err(ExpressionError::argument_count_error(1, args.len()));
450 }
451
452 let values = match &args[0] {
453 Value::List(list) => list,
454 _ => return Err(ExpressionError::type_error("First argument must be a list")),
455 };
456
457 if values.is_empty() {
458 return Ok(Value::string(String::new()));
459 }
460
461 let separator = match args.get(1) {
462 Some(Value::String(s)) => s.to_string(),
463 _ => ",".to_string(),
464 };
465
466 let mut indexed_values: Vec<(usize, &Value)> = values.iter().enumerate().collect();
467
468 if let Some(Value::List(order_by)) = args.get(2) {
469 if !order_by.values.is_empty() {
470 indexed_values.sort_by(|(_, val_a), (_, val_b)| {
471 let cmp_a = format!("{}", val_a);
472 let cmp_b = format!("{}", val_b);
473 cmp_a.cmp(&cmp_b)
474 });
475 }
476 }
477
478 let result: Vec<String> = indexed_values
479 .iter()
480 .map(|(_, v)| format!("{}", v))
481 .collect();
482 Ok(Value::string(result.join(&separator)))
483 }
484 AggregateFunction::VecSum => {
485 if args.is_empty() {
486 return Err(ExpressionError::argument_count_error(1, args.len()));
487 }
488
489 let values = match &args[0] {
491 Value::List(list) => list,
492 _ => {
493 return Err(ExpressionError::type_error(
494 "First argument must be a list of vectors",
495 ))
496 }
497 };
498
499 if values.is_empty() {
500 return Ok(Value::Null(NullType::NaN));
501 }
502
503 let mut sum_vec: Option<Vec<f32>> = None;
504 for val in values.iter() {
505 if let Value::Vector(v) = val {
506 let data = v.to_dense();
507 match &mut sum_vec {
508 Some(sum) => {
509 if sum.len() != data.len() {
510 return Err(ExpressionError::type_error(
511 "Vector dimensions must match",
512 ));
513 }
514 for (i, &val) in data.iter().enumerate() {
515 sum[i] += val;
516 }
517 }
518 None => sum_vec = Some(data.clone()),
519 }
520 } else {
521 return Err(ExpressionError::type_error("All elements must be vectors"));
522 }
523 }
524
525 match sum_vec {
526 Some(data) => Ok(Value::vector(data)),
527 None => Ok(Value::Null(NullType::NaN)),
528 }
529 }
530 AggregateFunction::PercentileCont => {
531 if args.len() < 2 {
533 return Err(ExpressionError::argument_count_error(2, args.len()));
534 }
535 let percentile = match &args[1] {
536 Value::Int(v) => *v as f64,
537 Value::BigInt(v) => *v as f64,
538 Value::Float(v) => *v as f64,
539 Value::Double(v) => *v,
540 _ => return Err(ExpressionError::type_error("Percentile must be a number")),
541 };
542 if !(0.0..=100.0).contains(&percentile) {
543 return Err(ExpressionError::new(
544 ExpressionErrorType::InvalidOperation,
545 "Percentile must be between 0 and 100",
546 ));
547 }
548 let values = match &args[0] {
549 Value::List(list) => list,
550 _ => return Err(ExpressionError::type_error("First argument must be a list")),
551 };
552 if values.is_empty() {
553 return Ok(Value::Null(linkrs_core::NullType::NaN));
554 }
555 let mut numeric_values = Vec::new();
556 for value in values.iter() {
557 match value {
558 Value::SmallInt(v) => numeric_values.push(*v as f64),
559 Value::Int(v) => numeric_values.push(*v as f64),
560 Value::BigInt(v) => numeric_values.push(*v as f64),
561 Value::Float(v) => numeric_values.push(*v as f64),
562 Value::Double(v) => numeric_values.push(*v),
563 _ => continue,
564 }
565 }
566 if numeric_values.is_empty() {
567 return Ok(Value::Null(linkrs_core::NullType::NaN));
568 }
569 numeric_values
570 .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
571 let index = (percentile / 100.0) * (numeric_values.len() - 1) as f64;
572 let lower_index = index.floor() as usize;
573 let upper_index = index.ceil() as usize;
574 if lower_index == upper_index {
575 Ok(Value::Double(numeric_values[lower_index]))
576 } else {
577 let lower_value = numeric_values[lower_index];
578 let upper_value = numeric_values[upper_index];
579 let weight = index - lower_index as f64;
580 let interpolated = lower_value + weight * (upper_value - lower_value);
581 Ok(Value::Double(interpolated))
582 }
583 }
584 AggregateFunction::StddevPop => {
585 if args.is_empty() {
586 return Err(ExpressionError::argument_count_error(1, args.len()));
587 }
588 let values = match &args[0] {
589 Value::List(list) => list,
590 _ => return Err(ExpressionError::type_error("First argument must be a list")),
591 };
592 if values.is_empty() {
593 return Ok(Value::Null(linkrs_core::NullType::NaN));
594 }
595 let mut numeric_values = Vec::new();
596 for value in values.iter() {
597 match value {
598 Value::SmallInt(v) => numeric_values.push(*v as f64),
599 Value::Int(v) => numeric_values.push(*v as f64),
600 Value::BigInt(v) => numeric_values.push(*v as f64),
601 Value::Float(v) => numeric_values.push(*v as f64),
602 Value::Double(v) => numeric_values.push(*v),
603 _ => continue,
604 }
605 }
606 if numeric_values.is_empty() {
607 return Ok(Value::Null(linkrs_core::NullType::NaN));
608 }
609 let n = numeric_values.len() as f64;
610 let mean: f64 = numeric_values.iter().sum::<f64>() / n;
611 let variance: f64 = numeric_values
612 .iter()
613 .map(|value| (value - mean).powi(2))
614 .sum::<f64>()
615 / n;
616 Ok(Value::Double(variance.sqrt()))
617 }
618 AggregateFunction::StddevSamp => {
619 if args.is_empty() {
620 return Err(ExpressionError::argument_count_error(1, args.len()));
621 }
622 let values = match &args[0] {
623 Value::List(list) => list,
624 _ => return Err(ExpressionError::type_error("First argument must be a list")),
625 };
626 if values.is_empty() {
627 return Ok(Value::Null(linkrs_core::NullType::NaN));
628 }
629 let mut numeric_values = Vec::new();
630 for value in values.iter() {
631 match value {
632 Value::SmallInt(v) => numeric_values.push(*v as f64),
633 Value::Int(v) => numeric_values.push(*v as f64),
634 Value::BigInt(v) => numeric_values.push(*v as f64),
635 Value::Float(v) => numeric_values.push(*v as f64),
636 Value::Double(v) => numeric_values.push(*v),
637 _ => continue,
638 }
639 }
640 if numeric_values.is_empty() {
641 return Ok(Value::Null(linkrs_core::NullType::NaN));
642 }
643 let n = numeric_values.len() as f64;
644 if n < 2.0 {
645 return Ok(Value::Null(linkrs_core::NullType::NaN));
646 }
647 let mean: f64 = numeric_values.iter().sum::<f64>() / n;
648 let variance: f64 = numeric_values
649 .iter()
650 .map(|value| (value - mean).powi(2))
651 .sum::<f64>()
652 / (n - 1.0);
653 Ok(Value::Double(variance.sqrt()))
654 }
655 AggregateFunction::Product => {
656 if args.is_empty() {
657 return Err(ExpressionError::argument_count_error(1, args.len()));
658 }
659 let values = match &args[0] {
660 Value::List(list) => list,
661 _ => return Err(ExpressionError::type_error("First argument must be a list")),
662 };
663 if values.is_empty() {
664 return Ok(Value::Null(linkrs_core::NullType::NaN));
665 }
666 let mut product = Value::Int(1);
667 for value in values.iter() {
668 if value.is_null() || value.is_empty() {
669 continue;
670 }
671 product = product.mul(value).map_err(ExpressionError::runtime_error)?;
672 }
673 Ok(product)
674 }
675 AggregateFunction::VecAvg => {
676 if args.is_empty() {
677 return Err(ExpressionError::argument_count_error(1, args.len()));
678 }
679
680 let values = match &args[0] {
682 Value::List(list) => list,
683 _ => {
684 return Err(ExpressionError::type_error(
685 "First argument must be a list of vectors",
686 ))
687 }
688 };
689
690 if values.is_empty() {
691 return Ok(Value::Null(NullType::NaN));
692 }
693
694 let count = values.len() as f32;
695 let mut sum_vec: Option<Vec<f32>> = None;
696 for val in values.iter() {
697 if let Value::Vector(v) = val {
698 let data = v.to_dense();
699 match &mut sum_vec {
700 Some(sum) => {
701 if sum.len() != data.len() {
702 return Err(ExpressionError::type_error(
703 "Vector dimensions must match",
704 ));
705 }
706 for (i, &val) in data.iter().enumerate() {
707 sum[i] += val;
708 }
709 }
710 None => sum_vec = Some(data.clone()),
711 }
712 } else {
713 return Err(ExpressionError::type_error("All elements must be vectors"));
714 }
715 }
716
717 match sum_vec {
718 Some(mut data) => {
719 for val in data.iter_mut() {
720 *val /= count;
721 }
722 Ok(Value::vector(data))
723 }
724 None => Ok(Value::Null(NullType::NaN)),
725 }
726 }
727 }
728 }
729}