1use std::cmp::PartialOrd;
2use std::collections::{BTreeMap, BTreeSet};
3use std::convert::TryFrom;
4use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
5use std::str::FromStr;
6
7use chrono::prelude::*;
8use chrono::serde::ts_seconds;
9use indexmap::IndexMap;
10use ordered_float::OrderedFloat;
11use rayon::prelude::*;
12use serde_derive::{Deserialize, Serialize};
13
14use crate::planner::{self, WhereCond};
15use crate::sql::Selector;
16use crate::sql::SelectorNode;
17use crate::value::PqlVector;
18
19#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
20#[serde(untagged)]
21pub enum BPqlValue {
22 Null,
23 Str(String),
24 Boolean(bool),
25 Float(OrderedFloat<f64>),
26 Int(i64),
27 #[serde(with = "ts_seconds")]
28 DateTime(DateTime<Utc>),
29 Array(BTreeSet<Self>),
30 Object(BTreeMap<String, Self>),
31}
32
33impl From<PqlValue> for BPqlValue {
34 fn from(pqlv: PqlValue) -> Self {
35 match pqlv {
36 PqlValue::Null => Self::Null,
37 PqlValue::Str(s) => Self::Str(s),
38 PqlValue::Boolean(b) => Self::Boolean(b),
39 PqlValue::Int(i) => Self::Int(i),
40 PqlValue::Float(f) => Self::Float(f),
41 PqlValue::DateTime(t) => Self::DateTime(t),
42 PqlValue::Array(_) => todo!(),
43 PqlValue::Object(_) => todo!(),
44 }
45 }
46}
47
48#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
49#[serde(untagged)]
50pub enum PqlValue {
51 Null,
52 Str(String),
53 Boolean(bool),
54 Int(i64),
55 Float(OrderedFloat<f64>),
56 #[serde(with = "ts_seconds")]
57 DateTime(DateTime<Utc>),
58 Array(Vec<Self>),
59 Object(IndexMap<String, Self>),
60}
61
62impl Default for PqlValue {
63 fn default() -> Self {
64 Self::Null
65 }
66}
67
68impl FromStr for PqlValue {
69 type Err = anyhow::Error;
70
71 fn from_str(s: &str) -> anyhow::Result<Self> {
72 crate::pqlir_parser::from_str(s)
73 }
74}
75
76impl From<&str> for PqlValue {
77 fn from(s: &str) -> Self {
78 Self::Str(s.to_owned())
79 }
80}
81
82impl From<bool> for PqlValue {
83 fn from(b: bool) -> Self {
84 Self::Boolean(b)
85 }
86}
87
88impl From<i64> for PqlValue {
89 fn from(i: i64) -> Self {
90 Self::Int(i)
91 }
92}
93
94impl From<f64> for PqlValue {
95 fn from(f: f64) -> Self {
96 Self::Float(OrderedFloat(f))
97 }
98}
99
100impl From<Vec<PqlValue>> for PqlValue {
101 fn from(v: Vec<PqlValue>) -> Self {
102 Self::Array(v)
103 }
104}
105
106impl PqlValue {
107 pub fn get(self, key: &str) -> Option<Self> {
108 match self {
109 Self::Object(map) => map.get(key).map(|v| v.to_owned()),
110 _ => None,
111 }
112 }
113
114 pub fn get_path(self, path: &[&str]) -> Option<Self> {
115 if let Some((key, path)) = path.split_first() {
116 if let Some(obj) = self.get(key) {
117 if path.is_empty() {
118 Some(obj)
119 } else {
120 obj.get_path(path)
121 }
122 } else {
123 None
124 }
125 } else {
126 unreachable!();
127 }
128 }
129
130 pub fn select_by_key(&self, key: &SelectorNode) -> Option<Self> {
131 match (self, key.to_owned()) {
132 (Self::Object(map), SelectorNode::String(key_s)) => {
133 map.get(&key_s).map(|v| v.to_owned())
134 }
135 _ => None,
136 }
137 }
138
139 pub fn get_mut_by_selectornode(&mut self, node: &SelectorNode) -> Option<&mut Self> {
140 match (self, node.to_owned()) {
141 (Self::Object(map), SelectorNode::String(key_s)) => map.get_mut(&key_s),
142 _ => None,
143 }
144 }
145
146 pub fn select_by_selector(&self, selector: &Selector) -> Option<Self> {
147 match self {
148 Self::Object(_map) => {
149 if let Some((key, tail)) = selector.split_first() {
150 if let Some(obj) = self.select_by_key(&key) {
151 obj.select_by_selector(&tail)
152 } else {
153 None
154 }
155 } else {
156 Some(self.to_owned())
157 }
158 }
159 Self::Array(array) => {
160 if let Some((key, _tail)) = selector.split_first() {
161 match key {
162 SelectorNode::Number(key_i) => {
163 if key_i < 0 {
164 todo!()
165 } else {
166 let key_u = key_i as usize;
167 array.get(key_u).map(|v| v.to_owned())
168 }
169 }
170 _ => {
171 let new_array = array
172 .into_iter()
173 .filter_map(|value| value.select_by_selector(&selector))
174 .collect::<Vec<_>>();
175 Some(Self::Array(new_array))
176 }
177 }
178 } else {
179 let new_array = array
180 .into_iter()
181 .filter_map(|value| value.select_by_selector(&selector))
182 .collect::<Vec<_>>();
183 Some(Self::Array(new_array))
184 }
185 }
186 _ => Some(self.clone()),
187 }
188 }
189
190 pub fn get_mut_by_selector(&mut self, selector: &Selector) -> Option<&mut Self> {
191 match self {
192 Self::Object(_map) => {
193 if let Some((key, tail)) = selector.split_first() {
194 if let Some(obj) = self.get_mut_by_selectornode(&key) {
195 obj.get_mut_by_selector(&tail)
196 } else {
197 None
198 }
199 } else {
200 Some(self)
201 }
202 }
203 Self::Array(array) => {
204 if let Some((key, _tail)) = selector.split_first() {
205 match key {
206 SelectorNode::Number(key_i) => {
207 if key_i < 0 {
208 todo!()
209 } else {
210 let key_u = key_i as usize;
211 array.get_mut(key_u)
212 }
213 }
214 _ => {
215 todo!()
216 }
217 }
218 } else {
219 todo!()
220 }
221 }
222 _ => Some(self),
223 }
224 }
225
226 pub fn print(&self) -> anyhow::Result<()> {
227 println!("{}", self.to_json()?);
228 Ok(())
229 }
230
231 pub fn to_json(&self) -> serde_json::Result<String> {
232 self.to_jsonp()
233 }
234
235 pub fn to_jsonp(&self) -> serde_json::Result<String> {
236 serde_json::to_string_pretty(self)
237 }
238
239 pub fn to_jsonc(&self) -> serde_json::Result<String> {
240 serde_json::to_string(self)
241 }
242
243 pub fn restrict(self, selector: &Selector, cond: &Option<WhereCond>) -> Option<Self> {
244 planner::filter::restrict(Some(self), selector, cond)
245 }
246}
247
248impl Neg for PqlValue {
249 type Output = Self;
250 fn neg(self) -> Self::Output {
251 match self {
252 Self::Int(a) => Self::Int(-a),
253 Self::Float(a) => Self::Float(-a),
254 _ => todo!(),
255 }
256 }
257}
258
259impl Add for PqlValue {
260 type Output = Self;
261 fn add(self, other: Self) -> Self::Output {
262 match (self, other) {
263 (Self::Int(a), Self::Int(b)) => Self::Int(a + b),
264 (Self::Int(a), Self::Float(b)) => Self::Float(OrderedFloat(a as f64) + b),
265 (Self::Float(a), Self::Int(b)) => Self::Float(a + OrderedFloat(b as f64)),
266 (Self::Float(a), Self::Float(b)) => Self::Float(a + b),
267 (Self::Array(array_a), Self::Array(array_b)) => {
268 let (vec_a, vec_b) = (PqlVector(array_a), PqlVector(array_b));
269 PqlValue::from(vec_a + vec_b)
270 }
271 (Self::Array(array), val) => {
272 let n = array.len();
273 let (vec_a, vec_b) = (PqlVector(array), PqlVector(vec![val; n]));
274 PqlValue::from(vec_a + vec_b)
275 }
276 (val, Self::Array(array)) => {
277 let n = array.len();
278 let (vec_a, vec_b) = (PqlVector(vec![val; n]), PqlVector(array));
279 PqlValue::from(vec_a + vec_b)
280 }
281 _ => todo!(),
282 }
283 }
284}
285
286impl Sub for PqlValue {
287 type Output = Self;
288 fn sub(self, other: Self) -> Self::Output {
289 match (self, other) {
290 (Self::Int(a), Self::Int(b)) => Self::Int(a - b),
291 (Self::Int(a), Self::Float(b)) => Self::Float(OrderedFloat(a as f64) - b),
292 (Self::Float(a), Self::Int(b)) => Self::Float(a - OrderedFloat(b as f64)),
293 (Self::Float(a), Self::Float(b)) => Self::Float(a - b),
294 (Self::Array(array_a), Self::Array(array_b)) => {
295 let (vec_a, vec_b) = (PqlVector(array_a), PqlVector(array_b));
296 PqlValue::from(vec_a - vec_b)
297 }
298 (Self::Array(array), val) => {
299 let n = array.len();
300 let (vec_a, vec_b) = (PqlVector(array), PqlVector(vec![val; n]));
301 PqlValue::from(vec_a - vec_b)
302 }
303 (val, Self::Array(array)) => {
304 let n = array.len();
305 let (vec_a, vec_b) = (PqlVector(vec![val; n]), PqlVector(array));
306 PqlValue::from(vec_a - vec_b)
307 }
308 _ => todo!(),
309 }
310 }
311}
312
313impl Mul for PqlValue {
314 type Output = Self;
315 fn mul(self, other: Self) -> Self::Output {
316 match (self.to_owned(), other.to_owned()) {
317 (Self::Int(a), Self::Int(b)) => Self::Int(a * b),
318 (Self::Int(a), Self::Float(b)) => Self::Float(OrderedFloat(a as f64) * b),
319 (Self::Float(a), Self::Int(b)) => Self::Float(a * OrderedFloat(b as f64)),
320 (Self::Float(a), Self::Float(b)) => Self::Float(a * b),
321 (Self::Array(array_a), Self::Array(array_b)) => {
322 let (vec_a, vec_b) = (PqlVector(array_a), PqlVector(array_b));
323 PqlValue::from(vec_a * vec_b)
324 }
325 (Self::Array(array), val) => {
326 let n = array.len();
327 let (vec_a, vec_b) = (PqlVector(array), PqlVector(vec![val; n]));
328 PqlValue::from(vec_a * vec_b)
329 }
330 (val, Self::Array(array)) => {
331 let n = array.len();
332 let (vec_a, vec_b) = (PqlVector(vec![val; n]), PqlVector(array));
333 PqlValue::from(vec_a * vec_b)
334 }
335 _ => todo!(),
336 }
337 }
338}
339
340impl Div for PqlValue {
341 type Output = Self;
342 fn div(self, other: Self) -> Self::Output {
343 match (self, other) {
344 (Self::Int(a), Self::Int(b)) => Self::Float(OrderedFloat(a as f64 / b as f64)),
345 (Self::Int(a), Self::Float(b)) => Self::Float(OrderedFloat(a as f64) / b),
346 (Self::Float(a), Self::Int(b)) => Self::Float(a / OrderedFloat(b as f64)),
347 (Self::Float(a), Self::Float(b)) => Self::Float(a / b),
348 (Self::Array(array_a), Self::Array(array_b)) => {
349 let (vec_a, vec_b) = (PqlVector(array_a), PqlVector(array_b));
350 PqlValue::from(vec_a / vec_b)
351 }
352 (Self::Array(array), val) => {
353 let n = array.len();
354 let (vec_a, vec_b) = (PqlVector(array), PqlVector(vec![val; n]));
355 PqlValue::from(vec_a / vec_b)
356 }
357 (val, Self::Array(array)) => {
358 let n = array.len();
359 let (vec_a, vec_b) = (PqlVector(vec![val; n]), PqlVector(array));
360 PqlValue::from(vec_a / vec_b)
361 }
362 _ => todo!(),
363 }
364 }
365}
366
367impl Rem for PqlValue {
368 type Output = Self;
369 fn rem(self, other: Self) -> Self::Output {
370 let (a, b) = match (self, other) {
371 (Self::Int(a), Self::Int(b)) => (a as f64, b as f64),
372 (Self::Int(a), Self::Float(OrderedFloat(b))) => (a as f64, b),
373 (Self::Float(OrderedFloat(a)), Self::Int(b)) => (a, b as f64),
374 (Self::Float(OrderedFloat(a)), Self::Float(OrderedFloat(b))) => (a, b),
375 _ => todo!(),
376 };
377 Self::from(a % b)
378 }
379}
380
381impl PqlValue {
382 pub fn powf(self, other: Self) -> Self {
383 let (a, b) = match (self, other) {
384 (Self::Int(a), Self::Int(b)) => (a as f64, b as f64),
385 (Self::Int(a), Self::Float(OrderedFloat(b))) => (a as f64, b),
386 (Self::Float(OrderedFloat(a)), Self::Int(b)) => (a, b as f64),
387 (Self::Float(OrderedFloat(a)), Self::Float(OrderedFloat(b))) => (a, b),
388 _ => todo!(),
389 };
390 Self::from(a.powf(b))
391 }
392}
393
394impl TryFrom<PqlValue> for i64 {
395 type Error = anyhow::Error;
396 fn try_from(value: PqlValue) -> anyhow::Result<Self> {
397 match value {
398 PqlValue::Int(int) => Ok(int),
399 PqlValue::Float(OrderedFloat(f)) => Ok(f as i64),
400 _ => anyhow::bail!("not numeric"),
401 }
402 }
403}
404
405impl PartialOrd for PqlValue {
406 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
407 let v1 = BPqlValue::from(self.to_owned());
408 let v2 = BPqlValue::from(other.to_owned());
409 Some(v1.cmp(&v2))
410 }
411}
412
413#[cfg(test)]
414mod tests {
415 use std::collections::VecDeque;
416 use std::str::FromStr;
417
418 use indexmap::IndexMap as Map;
419 use ordered_float::OrderedFloat;
420
421 use crate::parser;
422 use crate::planner::LogicalPlan;
423 use crate::pqlir_parser;
424 use crate::sql::Env;
425 use crate::sql::Expr;
426 use crate::sql::Selector;
427 use crate::sql::SelectorNode;
428 use crate::sql::Sql;
429 use crate::value::PqlValue;
430
431 #[test]
432 fn add_sub_mul_div() {
433 assert_eq!(
434 PqlValue::Float(OrderedFloat(1.)) + PqlValue::Float(OrderedFloat(2.)),
435 PqlValue::Float(OrderedFloat(3.))
436 );
437 assert_eq!(
438 PqlValue::Float(OrderedFloat(1.)) / PqlValue::Float(OrderedFloat(0.)),
439 PqlValue::Float(OrderedFloat(f64::INFINITY))
440 );
441 }
442
443 #[test]
444 fn select_at_arr_1() -> anyhow::Result<()> {
445 let value = PqlValue::from_str(r#"{ "arr" : [1,2,4] }"#)?;
446
447 let selected_value = value.select_by_selector(&Selector {
448 data: vec![
449 SelectorNode::String(String::from("arr")),
450 SelectorNode::Number(1),
451 ]
452 .into_iter()
453 .collect::<VecDeque<SelectorNode>>(),
454 });
455
456 assert_eq!(selected_value, Some(pqlir_parser::from_str("2")?));
457 Ok(())
458 }
459
460 #[test]
461 fn test_ord() {
462 let i1 = PqlValue::from(1);
463 let f2 = PqlValue::from(2.);
464 let i3 = PqlValue::from(3);
465
466 assert_eq!(i1 > f2, true);
467 assert_eq!(f2 < i3, true);
468 }
469
470 #[test]
471 fn test_update_value() -> anyhow::Result<()> {
472 let mut value = PqlValue::from_str(r#"{ "arr" : [1,2,4] }"#)?;
473
474 if let Some(partiql_value) = value.get_mut_by_selector(&Selector {
475 data: vec![
476 SelectorNode::String(String::from("arr")),
477 SelectorNode::Number(1),
478 ]
479 .into_iter()
480 .collect::<VecDeque<SelectorNode>>(),
481 }) {
482 *partiql_value = PqlValue::from(20.);
483 };
484 dbg!(&value);
485
486 assert_eq!(value, PqlValue::from_str(r#"{ "arr": [1,20,4] }"#)?);
487 Ok(())
488 }
489
490 #[test]
491 fn test_add() -> anyhow::Result<()> {
492 let data = PqlValue::from_str(
493 r#"
494{
495 "dat": [
496 { "n": 1 },
497 { "n": 2 },
498 { "n": 3 }
499 ]
500}
501"#,
502 )?;
503 let mut env = Env::default();
504 env.insert("", &Expr::from(data));
505
506 let mut sql = Sql::from_str(
507 r#"
508SELECT
509 dat.n + 3 AS n3,
510 4 + dat.n AS n4,
511 dat.n + dat.n AS nn,
512 "#,
513 )?;
514 let plan = LogicalPlan::from(sql);
515
516 let res = plan.execute(&mut env);
517
518 assert_eq!(
519 res,
520 PqlValue::from_str(
521 r#"
522[
523 {
524 "n3": 4.0,
525 "n4": 5.0,
526 "nn": 2.0
527 },
528 {
529 "n3": 5.0,
530 "n4": 6.0,
531 "nn": 4.0
532 },
533 {
534 "n3": 6.0,
535 "n4": 7.0,
536 "nn": 6.0
537 }
538]
539 "#
540 )?
541 );
542 Ok(())
543 }
544
545 #[test]
546 fn test_sub() -> anyhow::Result<()> {
547 let data = PqlValue::from_str(
548 r#"
549{
550 "dat": [
551 { "n": 1 },
552 { "n": 2 },
553 { "n": 3 }
554 ]
555}
556"#,
557 )?;
558 let mut env = Env::default();
559 env.insert("", &Expr::from(data));
560
561 let mut sql = Sql::from_str(
562 r#"
563SELECT
564 dat.n - 3 AS n3,
565 4 - dat.n AS n4,
566 dat.n - dat.n AS nn,
567 "#,
568 )?;
569 let plan = LogicalPlan::from(sql);
570
571 let res = plan.execute(&mut env);
572 res.print();
573
574 assert_eq!(
575 res,
576 PqlValue::from_str(
577 r#"
578[
579 {
580 "n3": -2.0,
581 "n4": 3.0,
582 "nn": 0.0
583 },
584 {
585 "n3": -1.0,
586 "n4": 2.0,
587 "nn": 0.0
588 },
589 {
590 "n3": 0.0,
591 "n4": 1.0,
592 "nn": 0.0
593 }
594]
595 "#
596 )?
597 );
598 Ok(())
599 }
600
601 #[test]
602 fn test_mul() -> anyhow::Result<()> {
603 let data = PqlValue::from_str(
604 r#"
605{
606 "dat": [
607 { "n": 1 },
608 { "n": 2 },
609 { "n": 3 }
610 ]
611}
612"#,
613 )?;
614 let mut env = Env::default();
615 env.insert("", &Expr::from(data));
616
617 let mut sql = Sql::from_str(
618 r#"
619SELECT
620 dat.n * 3 AS n3,
621 4* dat.n AS n4,
622 dat.n* dat.n AS nn,
623 "#,
624 )?;
625 let plan = LogicalPlan::from(sql);
626
627 let res = plan.execute(&mut env);
628
629 assert_eq!(
630 res,
631 PqlValue::from_str(
632 r#"
633[
634 {
635 "n3": 3.0,
636 "n4": 4.0,
637 "nn": 1.0
638 },
639 {
640 "n3": 6.0,
641 "n4": 8.0,
642 "nn": 4.0
643 },
644 {
645 "n3": 9.0,
646 "n4": 12.0,
647 "nn": 9.0
648 }
649]
650 "#
651 )?
652 );
653 Ok(())
654 }
655
656 #[test]
657 fn test_div() -> anyhow::Result<()> {
658 let data = PqlValue::from_str(
659 r#"
660{
661 "dat": [
662 { "n": 1 },
663 { "n": 2 },
664 { "n": 3 }
665 ]
666}
667"#,
668 )?;
669 let mut env = Env::default();
670 env.insert("", &Expr::from(data));
671
672 let mut sql = Sql::from_str(
673 r#"
674SELECT
675 dat.n / 3 AS n3,
676 4 / dat.n AS n4,
677 dat.n / dat.n AS nn,
678 "#,
679 )?;
680 let plan = LogicalPlan::from(sql);
681
682 let res = plan.execute(&mut env);
683 res.print();
684
685 assert_eq!(
686 res,
687 PqlValue::from_str(
688 r#"
689[
690 {
691 "n3": 0.3333333333333333,
692 "n4": 4.0,
693 "nn": 1.0
694 },
695 {
696 "n3": 0.6666666666666666,
697 "n4": 2.0,
698 "nn": 1.0
699 },
700 {
701 "n3": 1.0,
702 "n4": 1.3333333333333333,
703 "nn": 1.0
704 }
705]
706
707 "#
708 )?
709 );
710 Ok(())
711 }
712
713 #[test]
714 fn test_calc_bmi() -> anyhow::Result<()> {
715 let data = PqlValue::from_str(
716 r#"
717[
718 { "no": 1, "height": 0.7, "weight": 6.9 },
719 { "no": 2, "height": 1.0, "weight": 13.0 },
720 { "no": 3, "height": 2.0, "weight": 100.0 },
721 { "no": 4, "height": 0.6, "weight": 8.5 },
722 { "no": 5, "height": 1.1, "weight": 19.0 },
723 { "no": 6, "height": 1.7, "weight": 90.5 },
724 { "no": 7, "height": 0.5, "weight": 9.0 },
725 { "no": 8, "height": 1.0, "weight": 22.5 },
726 { "no": 9, "height": 1.6, "weight": 85.5 },
727 { "no": 10, "height": 0.3, "weight": 2.9 }
728]
729"#,
730 )?;
731 let mut env = Env::default();
732 env.insert("", &Expr::from(data));
733
734 let mut sql = Sql::from_str(
735 r#"
736SELECT
737 no,
738 weight/height/height AS bmi
739ORDER BY bmi DESC
740LIMIT 3
741 "#,
742 )?;
743 let plan = LogicalPlan::from(sql);
744
745 let res = plan.execute(&mut env);
746 dbg!(&res);
747 res.print();
748
749 assert_eq!(
750 res,
751 PqlValue::from_str(
752 r#"
753[
754 { "no": 7.0, "bmi": 36.0 },
755 { "no": 9.0, "bmi": 33.3984375 },
756 { "no": 10.0, "bmi": 32.22222222222222 }
757]
758 "#
759 )?
760 );
761
762 Ok(())
763 }
764}