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}