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