Skip to main content

linkrs_query/executor/expression/evaluator/
functions.rs

1//! Function call evaluation
2//!
3//! Provide the functionality to evaluate aggregate functions.
4
5use 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
11/// Function evaluator
12pub struct FunctionEvaluator;
13
14impl FunctionEvaluator {
15    /// Evaluation aggregate functions
16    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                // Sum all vectors in the list
490                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                // Same logic as PERCENTILE
532                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                // Average all vectors in the list
681                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}