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
806fn 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}