Skip to main content

teaql_runtime/
memory.rs

1use std::cmp::Ordering;
2use std::collections::BTreeMap;
3use std::sync::{Arc, Mutex};
4
5use rust_decimal::Decimal;
6use rust_decimal::prelude::ToPrimitive;
7use teaql_core::{
8    Aggregate, AggregateFunction, BinaryOp, DeleteCommand, Entity, Expr, ExprFunction,
9    InsertCommand, Record, RecoverCommand, RelationAggregate, SelectQuery, SmartList,
10    SortDirection, UpdateCommand, Value,
11};
12
13use crate::{DataServiceError, InMemoryMetadataStore, MetadataStore, RuntimeError};
14
15#[derive(Debug)]
16pub enum MemoryDataServiceError {
17    Poisoned,
18    UnsupportedExpression(String),
19    UnsupportedAggregate(String),
20}
21
22impl std::fmt::Display for MemoryDataServiceError {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        match self {
25            Self::Poisoned => write!(f, "memory data service lock poisoned"),
26            Self::UnsupportedExpression(message) => {
27                write!(f, "unsupported memory expression: {message}")
28            }
29            Self::UnsupportedAggregate(message) => {
30                write!(f, "unsupported memory aggregate: {message}")
31            }
32        }
33    }
34}
35
36impl std::error::Error for MemoryDataServiceError {}
37
38#[derive(Debug, Clone)]
39pub struct MemoryDataService<M = InMemoryMetadataStore> {
40    metadata: M,
41    data: Arc<Mutex<BTreeMap<String, Vec<Record>>>>,
42}
43
44impl<M> MemoryDataService<M>
45where
46    M: MetadataStore,
47{
48    pub fn new(metadata: M) -> Self {
49        Self {
50            metadata,
51            data: Arc::new(Mutex::new(BTreeMap::new())),
52        }
53    }
54
55    pub fn with_rows(mut self, entity: impl Into<String>, rows: Vec<Record>) -> Self {
56        self.seed(entity, rows);
57        self
58    }
59
60    pub fn seed(&mut self, entity: impl Into<String>, rows: Vec<Record>) {
61        if let Ok(mut data) = self.data.lock() {
62            data.insert(entity.into(), rows);
63        }
64    }
65
66    pub fn fetch_all(
67        &self,
68        query: &SelectQuery,
69    ) -> Result<Vec<Record>, DataServiceError<MemoryDataServiceError>> {
70        self.require_entity(&query.entity)?;
71        let data = self
72            .data
73            .lock()
74            .map_err(|_| DataServiceError::Executor(MemoryDataServiceError::Poisoned))?;
75        let mut rows = data.get(&query.entity).cloned().unwrap_or_default();
76        drop(data);
77
78        if let Some(filter) = &query.filter {
79            rows = rows
80                .into_iter()
81                .filter_map(|row| match eval_filter(filter, &row) {
82                    Ok(true) => Some(Ok(row)),
83                    Ok(false) => None,
84                    Err(err) => Some(Err(err)),
85                })
86                .collect::<Result<Vec<_>, _>>()
87                .map_err(DataServiceError::Executor)?;
88        }
89
90        if !query.aggregates.is_empty() {
91            return aggregate_rows(query, &rows).map_err(DataServiceError::Executor);
92        }
93
94        apply_ordering(&mut rows, query);
95        rows = apply_slice(rows, query);
96        if !query.projection.is_empty() || !query.expr_projection.is_empty() {
97            rows = rows
98                .into_iter()
99                .map(|row| project_row(row, query))
100                .collect::<Result<Vec<_>, _>>()
101                .map_err(DataServiceError::Executor)?;
102        }
103        Ok(rows)
104    }
105
106    pub fn fetch_smart_list(
107        &self,
108        query: &SelectQuery,
109    ) -> Result<SmartList<Record>, DataServiceError<MemoryDataServiceError>> {
110        self.fetch_all(query).map(SmartList::from)
111    }
112
113    pub fn fetch_entities<T>(
114        &self,
115        query: &SelectQuery,
116    ) -> Result<SmartList<T>, DataServiceError<MemoryDataServiceError>>
117    where
118        T: Entity,
119    {
120        self.fetch_all(query)?
121            .into_iter()
122            .map(T::from_record)
123            .collect::<Result<Vec<_>, _>>()
124            .map(SmartList::from)
125            .map_err(DataServiceError::Entity)
126    }
127
128    pub fn fetch_all_with_relation_aggregates(
129        &self,
130        query: &SelectQuery,
131        relation_aggregates: &[RelationAggregate],
132    ) -> Result<Vec<Record>, DataServiceError<MemoryDataServiceError>> {
133        let mut rows = self.fetch_all(query)?;
134        self.enhance_relation_aggregates(&query.entity, &mut rows, relation_aggregates)?;
135        Ok(rows)
136    }
137
138    pub fn fetch_smart_list_with_relation_aggregates(
139        &self,
140        query: &SelectQuery,
141        relation_aggregates: &[RelationAggregate],
142    ) -> Result<SmartList<Record>, DataServiceError<MemoryDataServiceError>> {
143        self.fetch_all_with_relation_aggregates(query, relation_aggregates)
144            .map(SmartList::from)
145    }
146
147    pub fn fetch_entities_with_relation_aggregates<T>(
148        &self,
149        query: &SelectQuery,
150        relation_aggregates: &[RelationAggregate],
151    ) -> Result<SmartList<T>, DataServiceError<MemoryDataServiceError>>
152    where
153        T: Entity,
154    {
155        self.fetch_all_with_relation_aggregates(query, relation_aggregates)?
156            .into_iter()
157            .map(T::from_record)
158            .collect::<Result<Vec<_>, _>>()
159            .map(SmartList::from)
160            .map_err(DataServiceError::Entity)
161    }
162
163    pub fn enhance_relation_aggregates(
164        &self,
165        parent_entity: &str,
166        parent_rows: &mut [Record],
167        relation_aggregates: &[RelationAggregate],
168    ) -> Result<(), DataServiceError<MemoryDataServiceError>> {
169        for aggregate in relation_aggregates {
170            self.enhance_relation_aggregate(parent_entity, parent_rows, aggregate)?;
171        }
172        Ok(())
173    }
174
175    fn enhance_relation_aggregate(
176        &self,
177        parent_entity: &str,
178        parent_rows: &mut [Record],
179        aggregate: &RelationAggregate,
180    ) -> Result<(), DataServiceError<MemoryDataServiceError>> {
181        let descriptor = self.metadata.entity(parent_entity).ok_or_else(|| {
182            DataServiceError::Runtime(RuntimeError::MissingEntity(parent_entity.to_owned()))
183        })?;
184
185        let relation = descriptor
186            .relation_by_name(&aggregate.relation_name)
187            .ok_or_else(|| {
188                DataServiceError::Runtime(RuntimeError::MissingRelation {
189                    entity: parent_entity.to_owned(),
190                    relation: aggregate.relation_name.clone(),
191                })
192            })?;
193
194        let ids = parent_rows
195            .iter()
196            .filter_map(|row| row.get(&relation.local_key).cloned())
197            .collect::<Vec<_>>();
198
199        if ids.is_empty() {
200            let value = if aggregate.single_result {
201                Value::U64(0)
202            } else {
203                Value::List(Vec::new())
204            };
205            for parent in parent_rows.iter_mut() {
206                parent.insert(aggregate.alias.clone(), value.clone());
207            }
208            return Ok(());
209        }
210
211        let mut query = aggregate.query.clone();
212        query.entity = relation.target_entity.clone();
213        query.projection.clear();
214        query.expr_projection.clear();
215        query.order_by.clear();
216        query.slice = None;
217        query.relations.clear();
218        if query.aggregates.is_empty() {
219            let alias = if aggregate.single_result {
220                aggregate.alias.clone()
221            } else {
222                "count".to_owned()
223            };
224            query = query.aggregate(Aggregate::count(alias));
225        }
226        if !query
227            .group_by
228            .iter()
229            .any(|field| field == &relation.foreign_key)
230        {
231            query = query.group_by(relation.foreign_key.clone());
232        }
233        query = query.and_filter(Expr::in_list(relation.foreign_key.clone(), ids));
234
235        let aggregate_rows = self.fetch_all(&query)?;
236
237        let mut buckets: BTreeMap<String, Vec<Record>> = BTreeMap::new();
238        for mut row in aggregate_rows {
239            if let Some(key) = row.remove(&relation.foreign_key) {
240                let bucket_key = local_graph_identity_key(&key);
241                buckets.entry(bucket_key).or_default().push(row);
242            }
243        }
244
245        for parent in parent_rows {
246            let value = parent
247                .get(&relation.local_key)
248                .and_then(|local_value| buckets.get(&local_graph_identity_key(local_value)))
249                .map(|rows| {
250                    if aggregate.single_result {
251                        rows.first()
252                            .map(|row| {
253                                if row.len() == 1 {
254                                    row.values().next().cloned().unwrap_or(Value::Null)
255                                } else {
256                                    Value::object(row.clone())
257                                }
258                            })
259                            .unwrap_or(Value::U64(0))
260                    } else {
261                        Value::List(rows.iter().cloned().map(Value::object).collect())
262                    }
263                })
264                .unwrap_or_else(|| {
265                    if aggregate.single_result {
266                        Value::U64(0)
267                    } else {
268                        Value::List(Vec::new())
269                    }
270                });
271            parent.insert(aggregate.alias.clone(), value);
272        }
273
274        Ok(())
275    }
276
277    pub fn insert(
278        &self,
279        command: &InsertCommand,
280    ) -> Result<u64, DataServiceError<MemoryDataServiceError>> {
281        self.require_entity(&command.entity)?;
282        let mut data = self
283            .data
284            .lock()
285            .map_err(|_| DataServiceError::Executor(MemoryDataServiceError::Poisoned))?;
286        data.entry(command.entity.clone())
287            .or_default()
288            .push(command.values.clone());
289        Ok(1)
290    }
291
292    pub fn update(
293        &self,
294        command: &UpdateCommand,
295    ) -> Result<u64, DataServiceError<MemoryDataServiceError>> {
296        let (id_property, version_property) = self.id_and_version_properties(&command.entity)?;
297        let mut data = self
298            .data
299            .lock()
300            .map_err(|_| DataServiceError::Executor(MemoryDataServiceError::Poisoned))?;
301        let rows = data.entry(command.entity.clone()).or_default();
302        let Some(row) = rows
303            .iter_mut()
304            .find(|row| row.get(id_property) == Some(&command.id))
305        else {
306            return self.maybe_optimistic_conflict(
307                command.expected_version,
308                &command.entity,
309                &command.id,
310            );
311        };
312
313        if let Some(expected) = command.expected_version {
314            if row.get(version_property) != Some(&Value::I64(expected)) {
315                println!(
316                    "OptimisticLockConflict in memory.rs update! entity={}, id={:?}, expected={}, existing={:?}",
317                    command.entity,
318                    command.id,
319                    expected,
320                    row.get(version_property)
321                );
322                return Err(DataServiceError::Runtime(
323                    RuntimeError::OptimisticLockConflict {
324                        entity: command.entity.clone(),
325                        id: format!("{:?}", command.id),
326                    },
327                ));
328            }
329            row.insert(version_property.to_owned(), Value::I64(expected + 1));
330        }
331
332        for (key, value) in &command.values {
333            row.insert(key.clone(), value.clone());
334        }
335        Ok(1)
336    }
337
338    pub fn delete(
339        &self,
340        command: &DeleteCommand,
341    ) -> Result<u64, DataServiceError<MemoryDataServiceError>> {
342        let (id_property, version_property) = self.id_and_version_properties(&command.entity)?;
343        let mut data = self
344            .data
345            .lock()
346            .map_err(|_| DataServiceError::Executor(MemoryDataServiceError::Poisoned))?;
347        let rows = data.entry(command.entity.clone()).or_default();
348        let Some(index) = rows
349            .iter()
350            .position(|row| row.get(id_property) == Some(&command.id))
351        else {
352            return self.maybe_optimistic_conflict(
353                command.expected_version,
354                &command.entity,
355                &command.id,
356            );
357        };
358
359        if let Some(expected_version) = command.expected_version {
360            if rows[index].get(version_property) != Some(&Value::I64(expected_version)) {
361                return Err(DataServiceError::Runtime(
362                    RuntimeError::OptimisticLockConflict {
363                        entity: command.entity.clone(),
364                        id: format!("{:?}", command.id),
365                    },
366                ));
367            }
368        }
369
370        if command.soft_delete {
371            let next_version = command
372                .expected_version
373                .or_else(|| read_i64(rows[index].get(version_property)))
374                .map(|version| -(version.abs() + 1))
375                .unwrap_or(-1);
376            rows[index].insert(version_property.to_owned(), Value::I64(next_version));
377        } else {
378            rows.remove(index);
379        }
380        Ok(1)
381    }
382
383    pub fn recover(
384        &self,
385        command: &RecoverCommand,
386    ) -> Result<u64, DataServiceError<MemoryDataServiceError>> {
387        let (id_property, version_property) = self.id_and_version_properties(&command.entity)?;
388        let mut data = self
389            .data
390            .lock()
391            .map_err(|_| DataServiceError::Executor(MemoryDataServiceError::Poisoned))?;
392        let rows = data.entry(command.entity.clone()).or_default();
393        let Some(row) = rows
394            .iter_mut()
395            .find(|row| row.get(id_property) == Some(&command.id))
396        else {
397            return Err(DataServiceError::Runtime(
398                RuntimeError::OptimisticLockConflict {
399                    entity: command.entity.clone(),
400                    id: format!("{:?}", command.id),
401                },
402            ));
403        };
404
405        if row.get(version_property) != Some(&Value::I64(command.expected_version)) {
406            return Err(DataServiceError::Runtime(
407                RuntimeError::OptimisticLockConflict {
408                    entity: command.entity.clone(),
409                    id: format!("{:?}", command.id),
410                },
411            ));
412        }
413
414        row.insert(
415            version_property.to_owned(),
416            Value::I64(command.expected_version.abs() + 1),
417        );
418        Ok(1)
419    }
420
421    fn require_entity(&self, entity: &str) -> Result<(), DataServiceError<MemoryDataServiceError>> {
422        self.metadata.entity(entity).map(|_| ()).ok_or_else(|| {
423            DataServiceError::Runtime(RuntimeError::MissingEntity(entity.to_owned()))
424        })
425    }
426
427    fn id_and_version_properties(
428        &self,
429        entity: &str,
430    ) -> Result<(&str, &str), DataServiceError<MemoryDataServiceError>> {
431        let descriptor = self.metadata.entity(entity).ok_or_else(|| {
432            DataServiceError::Runtime(RuntimeError::MissingEntity(entity.to_owned()))
433        })?;
434        let id = descriptor
435            .id_property()
436            .map(|property| property.name.as_str())
437            .unwrap_or("id");
438        let version = descriptor
439            .version_property()
440            .map(|property| property.name.as_str())
441            .unwrap_or("version");
442        Ok((id, version))
443    }
444
445    fn maybe_optimistic_conflict(
446        &self,
447        expected_version: Option<i64>,
448        entity: &str,
449        id: &Value,
450    ) -> Result<u64, DataServiceError<MemoryDataServiceError>> {
451        if expected_version.is_some() {
452            Err(DataServiceError::Runtime(
453                RuntimeError::OptimisticLockConflict {
454                    entity: entity.to_owned(),
455                    id: format!("{id:?}"),
456                },
457            ))
458        } else {
459            Ok(0)
460        }
461    }
462}
463
464fn eval_filter(expr: &Expr, row: &Record) -> Result<bool, MemoryDataServiceError> {
465    match expr {
466        Expr::Column(_) | Expr::Value(_) | Expr::Function { .. } => {
467            value_truthy(&eval_value(expr, row)?)
468        }
469        Expr::Binary { left, op, right } => {
470            let left = eval_value(left, row)?;
471            let right = eval_value(right, row)?;
472            eval_binary(&left, *op, &right)
473        }
474        Expr::SubQuery { .. } => Err(MemoryDataServiceError::UnsupportedExpression(
475            "subquery filters require a SQL executor".to_owned(),
476        )),
477        Expr::Between { expr, lower, upper } => {
478            let value = eval_value(expr, row)?;
479            let lower = eval_value(lower, row)?;
480            let upper = eval_value(upper, row)?;
481            Ok(compare_values(&value, &lower) != Some(Ordering::Less)
482                && compare_values(&value, &upper) != Some(Ordering::Greater))
483        }
484        Expr::IsNull(expr) => Ok(matches!(eval_value(expr, row)?, Value::Null)),
485        Expr::IsNotNull(expr) => Ok(!matches!(eval_value(expr, row)?, Value::Null)),
486        Expr::And(parts) => {
487            for part in parts {
488                if !eval_filter(part, row)? {
489                    return Ok(false);
490                }
491            }
492            Ok(true)
493        }
494        Expr::Or(parts) => {
495            for part in parts {
496                if eval_filter(part, row)? {
497                    return Ok(true);
498                }
499            }
500            Ok(false)
501        }
502        Expr::Not(expr) => Ok(!eval_filter(expr, row)?),
503    }
504}
505
506fn eval_value(expr: &Expr, row: &Record) -> Result<Value, MemoryDataServiceError> {
507    match expr {
508        Expr::Column(column) => Ok(row.get(column).cloned().unwrap_or(Value::Null)),
509        Expr::Value(value) => Ok(value.clone()),
510        Expr::Function { function, args } => eval_function(*function, args, row),
511        other => Err(MemoryDataServiceError::UnsupportedExpression(format!(
512            "cannot evaluate {other:?} as a scalar value"
513        ))),
514    }
515}
516
517fn eval_function(
518    function: ExprFunction,
519    args: &[Expr],
520    row: &Record,
521) -> Result<Value, MemoryDataServiceError> {
522    match function {
523        ExprFunction::Soundex => {
524            let [arg] = args else {
525                return Err(MemoryDataServiceError::UnsupportedExpression(
526                    "SOUNDEX expects exactly one argument".to_owned(),
527                ));
528            };
529            match eval_value(arg, row)? {
530                Value::Text(value) => Ok(Value::Text(soundex(&value))),
531                Value::Null => Ok(Value::Null),
532                other => Err(MemoryDataServiceError::UnsupportedExpression(format!(
533                    "SOUNDEX expects text, got {other:?}"
534                ))),
535            }
536        }
537        ExprFunction::Gbk => {
538            let [arg] = args else {
539                return Err(MemoryDataServiceError::UnsupportedExpression(
540                    "GBK expects exactly one argument".to_owned(),
541                ));
542            };
543            eval_value(arg, row)
544        }
545        other => Err(MemoryDataServiceError::UnsupportedExpression(format!(
546            "function {other:?} is only supported by SQL execution"
547        ))),
548    }
549}
550
551fn eval_binary(left: &Value, op: BinaryOp, right: &Value) -> Result<bool, MemoryDataServiceError> {
552    match op {
553        BinaryOp::Eq => Ok(values_equal(left, right)),
554        BinaryOp::Ne => Ok(!values_equal(left, right)),
555        BinaryOp::Gt => Ok(compare_values(left, right) == Some(Ordering::Greater)),
556        BinaryOp::Gte => Ok(matches!(
557            compare_values(left, right),
558            Some(Ordering::Greater | Ordering::Equal)
559        )),
560        BinaryOp::Lt => Ok(compare_values(left, right) == Some(Ordering::Less)),
561        BinaryOp::Lte => Ok(matches!(
562            compare_values(left, right),
563            Some(Ordering::Less | Ordering::Equal)
564        )),
565        BinaryOp::Like => match (left, right) {
566            (Value::Text(value), Value::Text(pattern)) => Ok(like_matches(value, pattern)),
567            _ => Ok(false),
568        },
569        BinaryOp::NotLike => match (left, right) {
570            (Value::Text(value), Value::Text(pattern)) => Ok(!like_matches(value, pattern)),
571            _ => Ok(true),
572        },
573        BinaryOp::In | BinaryOp::InLarge => match right {
574            Value::List(values) => Ok(values.iter().any(|value| values_equal(left, value))),
575            _ => Err(MemoryDataServiceError::UnsupportedExpression(
576                "IN expects a list value".to_owned(),
577            )),
578        },
579        BinaryOp::NotIn | BinaryOp::NotInLarge => match right {
580            Value::List(values) => Ok(!values.iter().any(|value| values_equal(left, value))),
581            _ => Err(MemoryDataServiceError::UnsupportedExpression(
582                "NOT IN expects a list value".to_owned(),
583            )),
584        },
585    }
586}
587
588fn value_truthy(value: &Value) -> Result<bool, MemoryDataServiceError> {
589    match value {
590        Value::Bool(value) => Ok(*value),
591        Value::Null => Ok(false),
592        other => Err(MemoryDataServiceError::UnsupportedExpression(format!(
593            "non-boolean expression result: {other:?}"
594        ))),
595    }
596}
597
598fn values_equal(left: &Value, right: &Value) -> bool {
599    match (left, right) {
600        (Value::I64(left), Value::U64(right)) if *left >= 0 => *left as u64 == *right,
601        (Value::U64(left), Value::I64(right)) if *right >= 0 => *left == *right as u64,
602        _ => left == right,
603    }
604}
605
606fn compare_values(left: &Value, right: &Value) -> Option<Ordering> {
607    match (left, right) {
608        (Value::I64(left), Value::I64(right)) => Some(left.cmp(right)),
609        (Value::U64(left), Value::U64(right)) => Some(left.cmp(right)),
610        (Value::I64(left), Value::U64(right)) if *left >= 0 => Some((*left as u64).cmp(right)),
611        (Value::U64(left), Value::I64(right)) if *right >= 0 => Some(left.cmp(&(*right as u64))),
612        (Value::F64(left), Value::F64(right)) => left.partial_cmp(right),
613        (Value::Decimal(left), Value::Decimal(right)) => Some(left.cmp(right)),
614        (Value::Text(left), Value::Text(right)) => Some(left.cmp(right)),
615        (Value::Date(left), Value::Date(right)) => Some(left.cmp(right)),
616        (Value::Timestamp(left), Value::Timestamp(right)) => Some(left.cmp(right)),
617        _ => None,
618    }
619}
620
621fn like_matches(value: &str, pattern: &str) -> bool {
622    if pattern == "%" {
623        return true;
624    }
625    match (pattern.strip_prefix('%'), pattern.strip_suffix('%')) {
626        (Some(inner), Some(_)) if pattern.len() >= 2 => value.contains(&inner[..inner.len() - 1]),
627        (Some(suffix), None) => value.ends_with(suffix),
628        (None, Some(prefix)) => value.starts_with(prefix),
629        _ => value == pattern,
630    }
631}
632
633fn soundex(value: &str) -> String {
634    let mut letters = value
635        .chars()
636        .filter(|ch| ch.is_ascii_alphabetic())
637        .map(|ch| ch.to_ascii_uppercase());
638    let Some(first) = letters.next() else {
639        return "0000".to_owned();
640    };
641
642    let mut output = String::with_capacity(4);
643    output.push(first);
644    let mut previous_code = soundex_code(first);
645
646    for ch in letters {
647        let code = soundex_code(ch);
648        if code != '0' && code != previous_code {
649            output.push(code);
650            if output.len() == 4 {
651                return output;
652            }
653        }
654        previous_code = code;
655    }
656
657    while output.len() < 4 {
658        output.push('0');
659    }
660    output
661}
662
663fn soundex_code(ch: char) -> char {
664    match ch {
665        'B' | 'F' | 'P' | 'V' => '1',
666        'C' | 'G' | 'J' | 'K' | 'Q' | 'S' | 'X' | 'Z' => '2',
667        'D' | 'T' => '3',
668        'L' => '4',
669        'M' | 'N' => '5',
670        'R' => '6',
671        _ => '0',
672    }
673}
674
675fn apply_ordering(rows: &mut [Record], query: &SelectQuery) {
676    for order in query.order_by.iter().rev() {
677        rows.sort_by(|left, right| {
678            let left_value = if let Some(expr) = &order.expr {
679                eval_value(expr, left).ok()
680            } else {
681                left.get(&order.field).cloned()
682            };
683            let right_value = if let Some(expr) = &order.expr {
684                eval_value(expr, right).ok()
685            } else {
686                right.get(&order.field).cloned()
687            };
688            let ordering = match (left_value.as_ref(), right_value.as_ref()) {
689                (Some(left), Some(right)) => compare_values(left, right).unwrap_or(Ordering::Equal),
690                (None, Some(_)) => Ordering::Less,
691                (Some(_), None) => Ordering::Greater,
692                (None, None) => Ordering::Equal,
693            };
694            match order.direction {
695                SortDirection::Asc => ordering,
696                SortDirection::Desc => ordering.reverse(),
697            }
698        });
699    }
700}
701
702fn apply_slice(rows: Vec<Record>, query: &SelectQuery) -> Vec<Record> {
703    let Some(slice) = query.slice else {
704        return rows;
705    };
706    let offset = usize::try_from(slice.offset).unwrap_or(usize::MAX);
707    let limit = slice
708        .limit
709        .and_then(|limit| usize::try_from(limit).ok())
710        .unwrap_or(usize::MAX);
711    rows.into_iter().skip(offset).take(limit).collect()
712}
713
714fn project_row(row: Record, query: &SelectQuery) -> Result<Record, MemoryDataServiceError> {
715    let mut output: Record = query
716        .projection
717        .iter()
718        .filter_map(|field| row.get(field).cloned().map(|value| (field.clone(), value)))
719        .collect();
720    for projection in &query.expr_projection {
721        output.insert(
722            projection.alias.clone(),
723            eval_value(&projection.expr, &row)?,
724        );
725    }
726    Ok(output)
727}
728
729fn aggregate_rows(
730    query: &SelectQuery,
731    rows: &[Record],
732) -> Result<Vec<Record>, MemoryDataServiceError> {
733    let mut groups: BTreeMap<Vec<String>, Vec<&Record>> = BTreeMap::new();
734    if query.group_by.is_empty() {
735        groups.insert(Vec::new(), rows.iter().collect());
736    } else {
737        for row in rows {
738            let key = query
739                .group_by
740                .iter()
741                .map(|field| row.get(field).map(value_key).unwrap_or_default())
742                .collect::<Vec<_>>();
743            groups.entry(key).or_default().push(row);
744        }
745    }
746
747    let rows = groups
748        .into_values()
749        .map(|rows| {
750            let mut output = Record::new();
751            if let Some(first) = rows.first() {
752                for field in &query.group_by {
753                    if let Some(value) = first.get(field) {
754                        output.insert(field.clone(), value.clone());
755                    }
756                }
757            }
758            for aggregate in &query.aggregates {
759                let value = match aggregate.function {
760                    AggregateFunction::Count => {
761                        if aggregate.field == "*" {
762                            Value::U64(rows.len() as u64)
763                        } else {
764                            Value::U64(
765                                rows.iter()
766                                    .filter(|row| {
767                                        !matches!(
768                                            row.get(&aggregate.field),
769                                            None | Some(Value::Null)
770                                        )
771                                    })
772                                    .count() as u64,
773                            )
774                        }
775                    }
776                    AggregateFunction::Sum => numeric_sum(&rows, &aggregate.field)?,
777                    AggregateFunction::Avg => numeric_avg(&rows, &aggregate.field)?,
778                    AggregateFunction::Min => min_max(&rows, &aggregate.field, false)?,
779                    AggregateFunction::Max => min_max(&rows, &aggregate.field, true)?,
780                    AggregateFunction::Stddev => numeric_stddev(&rows, &aggregate.field, true)?,
781                    AggregateFunction::StddevPop => numeric_stddev(&rows, &aggregate.field, false)?,
782                    AggregateFunction::VarSamp => numeric_variance(&rows, &aggregate.field, true)?,
783                    AggregateFunction::VarPop => numeric_variance(&rows, &aggregate.field, false)?,
784                    AggregateFunction::BitAnd => {
785                        bit_aggregate(&rows, &aggregate.field, BitOp::And)?
786                    }
787                    AggregateFunction::BitOr => bit_aggregate(&rows, &aggregate.field, BitOp::Or)?,
788                    AggregateFunction::BitXor => {
789                        bit_aggregate(&rows, &aggregate.field, BitOp::Xor)?
790                    }
791                };
792                output.insert(aggregate.alias.clone(), value);
793            }
794            for projection in &query.expr_projection {
795                output.insert(
796                    projection.alias.clone(),
797                    eval_value(&projection.expr, &output)?,
798                );
799            }
800            Ok(output)
801        })
802        .collect::<Result<Vec<_>, _>>()?;
803    if let Some(having) = &query.having {
804        rows.into_iter()
805            .filter_map(|row| match eval_filter(having, &row) {
806                Ok(true) => Some(Ok(row)),
807                Ok(false) => None,
808                Err(err) => Some(Err(err)),
809            })
810            .collect()
811    } else {
812        Ok(rows)
813    }
814}
815
816fn numeric_sum(rows: &[&Record], field: &str) -> Result<Value, MemoryDataServiceError> {
817    let mut decimal_sum = Decimal::ZERO;
818    let mut integer_sum: i128 = 0;
819    let mut saw_decimal = false;
820    for value in rows.iter().filter_map(|row| row.get(field)) {
821        match value {
822            Value::I64(value) => {
823                integer_sum += i128::from(*value);
824                decimal_sum += Decimal::from(*value);
825            }
826            Value::U64(value) => {
827                integer_sum += i128::from(*value);
828                decimal_sum += Decimal::from(*value);
829            }
830            Value::F64(value) => {
831                saw_decimal = true;
832                decimal_sum += decimal_from_f64(*value);
833            }
834            Value::Decimal(value) => {
835                saw_decimal = true;
836                decimal_sum += *value;
837            }
838            Value::Null => {}
839            other => {
840                return Err(MemoryDataServiceError::UnsupportedAggregate(format!(
841                    "SUM does not support {other:?}"
842                )));
843            }
844        }
845    }
846    if saw_decimal {
847        Ok(Value::Decimal(decimal_sum))
848    } else if integer_sum >= 0 {
849        Ok(Value::U64(integer_sum as u64))
850    } else {
851        Ok(Value::I64(integer_sum as i64))
852    }
853}
854
855fn numeric_avg(rows: &[&Record], field: &str) -> Result<Value, MemoryDataServiceError> {
856    let mut sum = Decimal::ZERO;
857    let mut count: u64 = 0;
858    for value in rows.iter().filter_map(|row| row.get(field)) {
859        match value {
860            Value::I64(value) => {
861                sum += Decimal::from(*value);
862                count += 1;
863            }
864            Value::U64(value) => {
865                sum += Decimal::from(*value);
866                count += 1;
867            }
868            Value::F64(value) => {
869                sum += decimal_from_f64(*value);
870                count += 1;
871            }
872            Value::Decimal(value) => {
873                sum += *value;
874                count += 1;
875            }
876            Value::Null => {}
877            other => {
878                return Err(MemoryDataServiceError::UnsupportedAggregate(format!(
879                    "AVG does not support {other:?}"
880                )));
881            }
882        }
883    }
884    Ok(if count == 0 {
885        Value::Null
886    } else {
887        Value::Decimal(sum / Decimal::from(count))
888    })
889}
890
891fn decimal_from_f64(value: f64) -> Decimal {
892    Decimal::from_f64_retain(value).unwrap_or(Decimal::ZERO)
893}
894
895fn numeric_values(rows: &[&Record], field: &str) -> Result<Vec<f64>, MemoryDataServiceError> {
896    rows.iter()
897        .filter_map(|row| row.get(field))
898        .filter(|value| !matches!(value, Value::Null))
899        .map(|value| match value {
900            Value::I64(value) => Ok(*value as f64),
901            Value::U64(value) => Ok(*value as f64),
902            Value::F64(value) => Ok(*value),
903            Value::Decimal(value) => value.to_f64().ok_or_else(|| {
904                MemoryDataServiceError::UnsupportedAggregate(format!(
905                    "cannot convert decimal {value} to f64 for statistical aggregate"
906                ))
907            }),
908            other => Err(MemoryDataServiceError::UnsupportedAggregate(format!(
909                "numeric aggregate does not support {other:?}"
910            ))),
911        })
912        .collect()
913}
914
915fn numeric_variance(
916    rows: &[&Record],
917    field: &str,
918    sample: bool,
919) -> Result<Value, MemoryDataServiceError> {
920    let values = numeric_values(rows, field)?;
921    let count = values.len();
922    if count == 0 || (sample && count < 2) {
923        return Ok(Value::Null);
924    }
925    let mean = values.iter().sum::<f64>() / count as f64;
926    let sum = values
927        .iter()
928        .map(|value| {
929            let diff = value - mean;
930            diff * diff
931        })
932        .sum::<f64>();
933    let denominator = if sample { count - 1 } else { count } as f64;
934    Ok(Value::Decimal(decimal_from_f64(sum / denominator)))
935}
936
937fn numeric_stddev(
938    rows: &[&Record],
939    field: &str,
940    sample: bool,
941) -> Result<Value, MemoryDataServiceError> {
942    Ok(match numeric_variance(rows, field, sample)? {
943        Value::Decimal(value) => {
944            Value::Decimal(decimal_from_f64(value.to_f64().unwrap_or(0.0).sqrt()))
945        }
946        Value::Null => Value::Null,
947        other => other,
948    })
949}
950
951#[derive(Debug, Clone, Copy)]
952enum BitOp {
953    And,
954    Or,
955    Xor,
956}
957
958fn bit_aggregate(
959    rows: &[&Record],
960    field: &str,
961    op: BitOp,
962) -> Result<Value, MemoryDataServiceError> {
963    let mut selected: Option<i64> = None;
964    for value in rows.iter().filter_map(|row| row.get(field)) {
965        let value = match value {
966            Value::I64(value) => *value,
967            Value::U64(value) => i64::try_from(*value).map_err(|_| {
968                MemoryDataServiceError::UnsupportedAggregate(format!(
969                    "BIT aggregate u64 {value} exceeds i64 range"
970                ))
971            })?,
972            Value::Null => continue,
973            other => {
974                return Err(MemoryDataServiceError::UnsupportedAggregate(format!(
975                    "BIT aggregate does not support {other:?}"
976                )));
977            }
978        };
979        selected = Some(match (selected, op) {
980            (None, _) => value,
981            (Some(current), BitOp::And) => current & value,
982            (Some(current), BitOp::Or) => current | value,
983            (Some(current), BitOp::Xor) => current ^ value,
984        });
985    }
986    Ok(selected.map(Value::I64).unwrap_or(Value::Null))
987}
988
989fn min_max(rows: &[&Record], field: &str, max: bool) -> Result<Value, MemoryDataServiceError> {
990    let mut selected: Option<Value> = None;
991    for value in rows.iter().filter_map(|row| row.get(field)) {
992        if matches!(value, Value::Null) {
993            continue;
994        }
995        match &selected {
996            None => selected = Some(value.clone()),
997            Some(current) => {
998                let Some(ordering) = compare_values(value, current) else {
999                    return Err(MemoryDataServiceError::UnsupportedAggregate(format!(
1000                        "MIN/MAX does not support {value:?}"
1001                    )));
1002                };
1003                if (max && ordering == Ordering::Greater) || (!max && ordering == Ordering::Less) {
1004                    selected = Some(value.clone());
1005                }
1006            }
1007        }
1008    }
1009    Ok(selected.unwrap_or(Value::Null))
1010}
1011
1012fn value_key(value: &Value) -> String {
1013    match value {
1014        Value::Null => "null".to_owned(),
1015        Value::Bool(value) => format!("b:{value}"),
1016        Value::I64(value) => format!("i:{value}"),
1017        Value::U64(value) => format!("u:{value}"),
1018        Value::F64(value) => format!("f:{value}"),
1019        Value::Decimal(value) => format!("d:{value}"),
1020        Value::Text(value) => format!("t:{value}"),
1021        Value::Json(value) => format!("j:{value}"),
1022        Value::Date(value) => format!("d:{value}"),
1023        Value::Timestamp(value) => format!("ts:{}", value.to_rfc3339()),
1024        Value::Object(_) => "object".to_owned(),
1025        Value::List(_) => "list".to_owned(),
1026    }
1027}
1028
1029fn read_i64(value: Option<&Value>) -> Option<i64> {
1030    match value {
1031        Some(Value::I64(value)) => Some(*value),
1032        _ => None,
1033    }
1034}
1035
1036fn local_graph_identity_key(value: &Value) -> String {
1037    match value {
1038        Value::I64(val) if *val >= 0 => format!("u:{}", *val as u64),
1039        Value::U64(val) => format!("u:{val}"),
1040        Value::Null => "null".to_owned(),
1041        Value::Bool(v) => format!("b:{v}"),
1042        Value::I64(v) => format!("i:{v}"),
1043        Value::F64(v) => format!("f:{v}"),
1044        Value::Decimal(v) => format!("d:{v}"),
1045        Value::Text(v) => format!("t:{v}"),
1046        Value::Json(v) => format!("j:{v}"),
1047        Value::Date(v) => format!("d:{v}"),
1048        Value::Timestamp(v) => format!("ts:{}", v.to_rfc3339()),
1049        Value::Object(_) => "o".to_owned(),
1050        Value::List(_) => "l".to_owned(),
1051    }
1052}