1use std::cmp::Ordering;
21use std::fmt::{Display, Formatter};
22use std::str::FromStr;
23
24use chrono::Datelike;
25use fnv::FnvHashSet;
26use serde::{Deserialize, Deserializer, Serialize, Serializer};
27
28use super::values::decimal_utils::decimal_from_i128_with_scale;
29use super::values::temporal::date;
30use super::{Datum, PrimitiveLiteral};
31use crate::ErrorKind;
32use crate::error::{Error, Result, invalid_data};
33use crate::expr::{
34 BinaryExpression, BoundPredicate, BoundReference, Predicate, PredicateOperator, Reference,
35 SetExpression, UnaryExpression,
36};
37use crate::spec::Literal;
38use crate::spec::datatypes::{PrimitiveType, Type};
39use crate::transform::{BoxedTransformFunction, create_transform_function};
40
41const UNIX_EPOCH_YEAR: i32 = 1970;
43
44#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
57pub enum Transform {
58 Identity,
63 Bucket(u32),
81 Truncate(u32),
110 Year,
115 Month,
120 Day,
125 Hour,
130 Void,
139 Unknown,
141}
142
143impl Transform {
144 pub fn to_human_string(&self, field_type: &Type, value: Option<&Literal>) -> String {
207 let Some(value) = value.and_then(Literal::as_primitive_literal) else {
208 return "null".to_string();
209 };
210
211 match (*self, value) {
212 (Self::Void, _) => "null".to_string(),
213 (Self::Year, PrimitiveLiteral::Int(ordinal)) => Self::human_year(ordinal),
220 (Self::Month, PrimitiveLiteral::Int(ordinal)) => Self::human_month(ordinal),
221 (Self::Day, PrimitiveLiteral::Int(ordinal)) => Self::human_day(ordinal),
222 (Self::Hour, PrimitiveLiteral::Int(ordinal)) => Self::human_hour(ordinal),
223 (_, value) => {
224 let field_type = field_type.as_primitive_type().unwrap();
225 Datum::new(field_type.clone(), value).to_human_string()
226 }
227 }
228 }
229
230 fn human_year(year_ordinal: i32) -> String {
234 format!("{:04}", UNIX_EPOCH_YEAR + year_ordinal)
235 }
236
237 fn human_month(month_ordinal: i32) -> String {
245 format!(
246 "{:04}-{:02}",
247 UNIX_EPOCH_YEAR + month_ordinal.div_euclid(12),
248 1 + month_ordinal.rem_euclid(12)
249 )
250 }
251
252 fn human_day(day_ordinal: i32) -> String {
259 let date = date::days_to_date(day_ordinal);
260 format!("{:04}-{:02}-{:02}", date.year(), date.month(), date.day())
261 }
262
263 fn human_hour(hour_ordinal: i32) -> String {
270 format!(
271 "{}-{:02}",
272 Self::human_day(hour_ordinal.div_euclid(24)),
273 hour_ordinal.rem_euclid(24)
274 )
275 }
276
277 pub fn result_type(&self, input_type: &Type) -> Result<Type> {
280 match self {
281 Transform::Identity => {
282 if matches!(input_type, Type::Primitive(_)) {
283 Ok(input_type.clone())
284 } else {
285 Err(invalid_data!(
286 "{input_type} is not a valid input type of identity transform",
287 ))
288 }
289 }
290 Transform::Void => Ok(input_type.clone()),
291 Transform::Unknown => Ok(Type::Primitive(PrimitiveType::String)),
292 Transform::Bucket(_) => {
293 if let Type::Primitive(p) = input_type {
294 match p {
295 PrimitiveType::Int
296 | PrimitiveType::Long
297 | PrimitiveType::Decimal { .. }
298 | PrimitiveType::Date
299 | PrimitiveType::Time
300 | PrimitiveType::Timestamp
301 | PrimitiveType::Timestamptz
302 | PrimitiveType::TimestampNs
303 | PrimitiveType::TimestamptzNs
304 | PrimitiveType::String
305 | PrimitiveType::Uuid
306 | PrimitiveType::Fixed(_)
307 | PrimitiveType::Binary => Ok(Type::Primitive(PrimitiveType::Int)),
308 _ => Err(invalid_data!(
309 "{input_type} is not a valid input type of bucket transform",
310 )),
311 }
312 } else {
313 Err(invalid_data!(
314 "{input_type} is not a valid input type of bucket transform",
315 ))
316 }
317 }
318 Transform::Truncate(_) => {
319 if let Type::Primitive(p) = input_type {
320 match p {
321 PrimitiveType::Int
322 | PrimitiveType::Long
323 | PrimitiveType::String
324 | PrimitiveType::Binary
325 | PrimitiveType::Decimal { .. } => Ok(input_type.clone()),
326 _ => Err(invalid_data!(
327 "{input_type} is not a valid input type of truncate transform",
328 )),
329 }
330 } else {
331 Err(invalid_data!(
332 "{input_type} is not a valid input type of truncate transform",
333 ))
334 }
335 }
336 Transform::Year | Transform::Month => {
337 if let Type::Primitive(p) = input_type {
338 match p {
339 PrimitiveType::Timestamp
340 | PrimitiveType::Timestamptz
341 | PrimitiveType::TimestampNs
342 | PrimitiveType::TimestamptzNs
343 | PrimitiveType::Date => Ok(Type::Primitive(PrimitiveType::Int)),
344 _ => Err(invalid_data!(
345 "{input_type} is not a valid input type of {self} transform",
346 )),
347 }
348 } else {
349 Err(invalid_data!(
350 "{input_type} is not a valid input type of {self} transform",
351 ))
352 }
353 }
354 Transform::Day => {
355 if let Type::Primitive(p) = input_type {
356 match p {
357 PrimitiveType::Timestamp
358 | PrimitiveType::Timestamptz
359 | PrimitiveType::TimestampNs
360 | PrimitiveType::TimestamptzNs
361 | PrimitiveType::Date => Ok(Type::Primitive(PrimitiveType::Date)),
362 _ => Err(invalid_data!(
363 "{input_type} is not a valid input type of {self} transform",
364 )),
365 }
366 } else {
367 Err(invalid_data!(
368 "{input_type} is not a valid input type of {self} transform",
369 ))
370 }
371 }
372 Transform::Hour => {
373 if let Type::Primitive(p) = input_type {
374 match p {
375 PrimitiveType::Timestamp
376 | PrimitiveType::Timestamptz
377 | PrimitiveType::TimestampNs
378 | PrimitiveType::TimestamptzNs => Ok(Type::Primitive(PrimitiveType::Int)),
379 _ => Err(invalid_data!(
380 "{input_type} is not a valid input type of {self} transform",
381 )),
382 }
383 } else {
384 Err(invalid_data!(
385 "{input_type} is not a valid input type of {self} transform",
386 ))
387 }
388 }
389 }
390 }
391
392 pub fn preserves_order(&self) -> bool {
394 !matches!(
395 self,
396 Transform::Void | Transform::Bucket(_) | Transform::Unknown
397 )
398 }
399
400 pub fn dedup_name(&self) -> String {
403 match self {
404 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
405 "time".to_string()
406 }
407 _ => format!("{self}"),
408 }
409 }
410
411 pub fn satisfies_order_of(&self, other: &Self) -> bool {
417 match self {
418 Transform::Identity => other.preserves_order(),
419 Transform::Hour => matches!(
420 other,
421 Transform::Hour | Transform::Day | Transform::Month | Transform::Year
422 ),
423 Transform::Day => matches!(other, Transform::Day | Transform::Month | Transform::Year),
424 Transform::Month => matches!(other, Transform::Month | Transform::Year),
425 _ => self == other,
426 }
427 }
428
429 pub fn strict_project(
442 &self,
443 name: &str,
444 predicate: &BoundPredicate,
445 ) -> Result<Option<Predicate>> {
446 let func = create_transform_function(self)?;
447
448 match self {
449 Transform::Identity => match predicate {
450 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
451 BoundPredicate::Binary(expr) => Ok(Some(Predicate::Binary(BinaryExpression::new(
452 expr.op(),
453 Reference::new(name),
454 expr.literal().to_owned(),
455 )))),
456 BoundPredicate::Set(expr) => Ok(Some(Predicate::Set(SetExpression::new(
457 expr.op(),
458 Reference::new(name),
459 expr.literals().to_owned(),
460 )))),
461 _ => Ok(None),
462 },
463 Transform::Bucket(_) => match predicate {
464 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
465 BoundPredicate::Binary(expr) => {
466 self.project_binary_expr(name, PredicateOperator::NotEq, expr, &func)
467 }
468 BoundPredicate::Set(expr) => {
469 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
470 }
471 _ => Ok(None),
472 },
473 Transform::Truncate(width) => match predicate {
474 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
475 BoundPredicate::Binary(expr) => {
476 if matches!(
477 expr.term().field().field_type.as_primitive_type(),
478 Some(&PrimitiveType::Int)
479 | Some(&PrimitiveType::Long)
480 | Some(&PrimitiveType::Decimal { .. })
481 ) {
482 self.truncate_number_strict(name, expr, &func)
483 } else if expr.op() == PredicateOperator::StartsWith {
484 let len = match expr.literal().literal() {
485 PrimitiveLiteral::String(s) => s.len(),
486 PrimitiveLiteral::Binary(b) => b.len(),
487 _ => {
488 return Err(invalid_data!(
489 "Expected a string or binary literal, got: {:?}",
490 expr.literal()
491 ));
492 }
493 };
494 match len.cmp(&(*width as usize)) {
495 Ordering::Less => Ok(Some(Predicate::Binary(BinaryExpression::new(
496 PredicateOperator::StartsWith,
497 Reference::new(name),
498 expr.literal().to_owned(),
499 )))),
500 Ordering::Equal => Ok(Some(Predicate::Binary(BinaryExpression::new(
501 PredicateOperator::Eq,
502 Reference::new(name),
503 expr.literal().to_owned(),
504 )))),
505 Ordering::Greater => Ok(None),
506 }
507 } else if expr.op() == PredicateOperator::NotStartsWith {
508 let len = match expr.literal().literal() {
509 PrimitiveLiteral::String(s) => s.len(),
510 PrimitiveLiteral::Binary(b) => b.len(),
511 _ => {
512 return Err(invalid_data!(
513 "Expected a string or binary literal, got: {:?}",
514 expr.literal()
515 ));
516 }
517 };
518 match len.cmp(&(*width as usize)) {
519 Ordering::Less => Ok(Some(Predicate::Binary(BinaryExpression::new(
520 PredicateOperator::NotStartsWith,
521 Reference::new(name),
522 expr.literal().to_owned(),
523 )))),
524 Ordering::Equal => Ok(Some(Predicate::Binary(BinaryExpression::new(
525 PredicateOperator::NotEq,
526 Reference::new(name),
527 expr.literal().to_owned(),
528 )))),
529 Ordering::Greater => {
530 Ok(Some(Predicate::Binary(BinaryExpression::new(
531 expr.op(),
532 Reference::new(name),
533 func.transform_literal_result(expr.literal())?,
534 ))))
535 }
536 }
537 } else {
538 self.truncate_array_strict(name, expr, &func)
539 }
540 }
541 BoundPredicate::Set(expr) => {
542 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
543 }
544 _ => Ok(None),
545 },
546 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
547 match predicate {
548 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
549 BoundPredicate::Binary(expr) => self.truncate_number_strict(name, expr, &func),
550 BoundPredicate::Set(expr) => {
551 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
552 }
553 _ => Ok(None),
554 }
555 }
556 _ => Ok(None),
557 }
558 }
559
560 pub fn project(&self, name: &str, predicate: &BoundPredicate) -> Result<Option<Predicate>> {
573 let func = create_transform_function(self)?;
574
575 match self {
576 Transform::Identity => match predicate {
577 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
578 BoundPredicate::Binary(expr) => Ok(Some(Predicate::Binary(BinaryExpression::new(
579 expr.op(),
580 Reference::new(name),
581 expr.literal().to_owned(),
582 )))),
583 BoundPredicate::Set(expr) => Ok(Some(Predicate::Set(SetExpression::new(
584 expr.op(),
585 Reference::new(name),
586 expr.literals().to_owned(),
587 )))),
588 _ => Ok(None),
589 },
590 Transform::Bucket(_) => match predicate {
591 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
592 BoundPredicate::Binary(expr) => {
593 self.project_binary_expr(name, PredicateOperator::Eq, expr, &func)
594 }
595 BoundPredicate::Set(expr) => {
596 self.project_set_expr(expr, PredicateOperator::In, name, &func)
597 }
598 _ => Ok(None),
599 },
600 Transform::Truncate(width) => match predicate {
601 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
602 BoundPredicate::Binary(expr) => {
603 self.project_binary_with_adjusted_boundary(name, expr, &func, Some(*width))
604 }
605 BoundPredicate::Set(expr) => {
606 self.project_set_expr(expr, PredicateOperator::In, name, &func)
607 }
608 _ => Ok(None),
609 },
610 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
611 match predicate {
612 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
613 BoundPredicate::Binary(expr) => {
614 self.project_binary_with_adjusted_boundary(name, expr, &func, None)
615 }
616 BoundPredicate::Set(expr) => {
617 self.project_set_expr(expr, PredicateOperator::In, name, &func)
618 }
619 _ => Ok(None),
620 }
621 }
622 _ => Ok(None),
623 }
624 }
625
626 fn can_transform(&self, datum: &Datum) -> bool {
628 let input_type = datum.data_type().clone();
629 self.result_type(&Type::Primitive(input_type)).is_ok()
630 }
631
632 fn project_unary(op: PredicateOperator, name: &str) -> Result<Option<Predicate>> {
634 Ok(Some(Predicate::Unary(UnaryExpression::new(
635 op,
636 Reference::new(name),
637 ))))
638 }
639
640 fn project_binary_expr(
647 &self,
648 name: &str,
649 op: PredicateOperator,
650 expr: &BinaryExpression<BoundReference>,
651 func: &BoxedTransformFunction,
652 ) -> Result<Option<Predicate>> {
653 if expr.op() != op || !self.can_transform(expr.literal()) {
654 return Ok(None);
655 }
656
657 Ok(Some(Predicate::Binary(BinaryExpression::new(
658 expr.op(),
659 Reference::new(name),
660 func.transform_literal_result(expr.literal())?,
661 ))))
662 }
663
664 fn project_binary_with_adjusted_boundary(
673 &self,
674 name: &str,
675 expr: &BinaryExpression<BoundReference>,
676 func: &BoxedTransformFunction,
677 width: Option<u32>,
678 ) -> Result<Option<Predicate>> {
679 if !self.can_transform(expr.literal()) {
680 return Ok(None);
681 }
682
683 let op = &expr.op();
684 let datum = &expr.literal();
685
686 if let Some(boundary) = Self::adjust_boundary(op, datum)? {
687 let transformed_projection = func.transform_literal_result(&boundary)?;
688
689 let adjusted_projection =
690 self.adjust_time_projection(op, datum, &transformed_projection);
691
692 let adjusted_operator = Self::adjust_operator(op, datum, width);
693
694 if let Some(op) = adjusted_operator {
695 let predicate = match adjusted_projection {
696 None => Predicate::Binary(BinaryExpression::new(
697 op,
698 Reference::new(name),
699 transformed_projection,
700 )),
701 Some(AdjustedProjection::Single(d)) => {
702 Predicate::Binary(BinaryExpression::new(op, Reference::new(name), d))
703 }
704 Some(AdjustedProjection::Set(d)) => Predicate::Set(SetExpression::new(
705 PredicateOperator::In,
706 Reference::new(name),
707 d,
708 )),
709 };
710 return Ok(Some(predicate));
711 }
712 };
713
714 Ok(None)
715 }
716
717 fn project_set_expr(
720 &self,
721 expr: &SetExpression<BoundReference>,
722 op: PredicateOperator,
723 name: &str,
724 func: &BoxedTransformFunction,
725 ) -> Result<Option<Predicate>> {
726 if expr.op() != op || expr.literals().iter().any(|d| !self.can_transform(d)) {
727 return Ok(None);
728 }
729
730 let mut new_set = FnvHashSet::default();
731
732 for lit in expr.literals() {
733 let datum = func.transform_literal_result(lit)?;
734
735 if let Some(AdjustedProjection::Single(d)) =
736 self.adjust_time_projection(&op, lit, &datum)
737 {
738 new_set.insert(d);
739 };
740
741 new_set.insert(datum);
742 }
743
744 Ok(Some(Predicate::Set(SetExpression::new(
745 expr.op(),
746 Reference::new(name),
747 new_set,
748 ))))
749 }
750
751 fn adjust_boundary(op: &PredicateOperator, datum: &Datum) -> Result<Option<Datum>> {
761 let adjusted_boundary = match op {
762 PredicateOperator::LessThan => match (datum.data_type(), datum.literal()) {
763 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Some(Datum::int(v - 1)),
764 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Some(Datum::long(v - 1)),
765 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
766 Some(Datum::decimal(decimal_from_i128_with_scale(v - 1, 0))?)
767 }
768 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Some(Datum::date(v - 1)),
769 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
770 Some(Datum::timestamp_micros(v - 1))
771 }
772 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
773 Some(Datum::timestamptz_micros(v - 1))
774 }
775 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
776 Some(Datum::timestamp_nanos(v - 1))
777 }
778 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
779 Some(Datum::timestamptz_nanos(v - 1))
780 }
781 _ => Some(datum.to_owned()),
782 },
783 PredicateOperator::GreaterThan => match (datum.data_type(), datum.literal()) {
784 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Some(Datum::int(v + 1)),
785 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Some(Datum::long(v + 1)),
786 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
787 Some(Datum::decimal(decimal_from_i128_with_scale(v + 1, 0))?)
788 }
789 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Some(Datum::date(v + 1)),
790 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
791 Some(Datum::timestamp_micros(v + 1))
792 }
793 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
794 Some(Datum::timestamptz_micros(v + 1))
795 }
796 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
797 Some(Datum::timestamp_nanos(v + 1))
798 }
799 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
800 Some(Datum::timestamptz_nanos(v + 1))
801 }
802 _ => Some(datum.to_owned()),
803 },
804 PredicateOperator::Eq
805 | PredicateOperator::LessThanOrEq
806 | PredicateOperator::GreaterThanOrEq
807 | PredicateOperator::StartsWith
808 | PredicateOperator::NotStartsWith => Some(datum.to_owned()),
809 _ => None,
810 };
811
812 Ok(adjusted_boundary)
813 }
814
815 fn adjust_operator(
826 op: &PredicateOperator,
827 datum: &Datum,
828 width: Option<u32>,
829 ) -> Option<PredicateOperator> {
830 match op {
831 PredicateOperator::LessThan => Some(PredicateOperator::LessThanOrEq),
832 PredicateOperator::GreaterThan => Some(PredicateOperator::GreaterThanOrEq),
833 PredicateOperator::StartsWith => match datum.literal() {
834 PrimitiveLiteral::String(s) => {
835 if let Some(w) = width
836 && s.len() == w as usize
837 {
838 return Some(PredicateOperator::Eq);
839 };
840 Some(*op)
841 }
842 _ => Some(*op),
843 },
844 PredicateOperator::NotStartsWith => match datum.literal() {
845 PrimitiveLiteral::String(s) => {
846 if let Some(w) = width {
847 let w = w as usize;
848
849 if s.len() == w {
850 return Some(PredicateOperator::NotEq);
851 }
852
853 if s.len() < w {
854 return Some(*op);
855 }
856
857 return None;
858 };
859 Some(*op)
860 }
861 _ => Some(*op),
862 },
863 _ => Some(*op),
864 }
865 }
866
867 fn adjust_time_projection(
870 &self,
871 op: &PredicateOperator,
872 original: &Datum,
873 transformed: &Datum,
874 ) -> Option<AdjustedProjection> {
875 let should_adjust = match self {
876 Transform::Day => matches!(
877 original.data_type(),
878 PrimitiveType::Timestamp
879 | PrimitiveType::Timestamptz
880 | PrimitiveType::TimestampNs
881 | PrimitiveType::TimestamptzNs
882 ),
883 Transform::Year | Transform::Month => true,
884 _ => false,
885 };
886
887 if should_adjust && let &PrimitiveLiteral::Int(v) = transformed.literal() {
888 match op {
889 PredicateOperator::LessThan
890 | PredicateOperator::LessThanOrEq
891 | PredicateOperator::In => {
892 if v < 0 {
893 match self {
896 Transform::Day => {
897 return Some(AdjustedProjection::Single(Datum::date(v + 1)));
898 }
899 _ => {
900 return Some(AdjustedProjection::Single(Datum::int(v + 1)));
901 }
902 }
903 };
904 }
905 PredicateOperator::Eq => {
906 if v < 0 {
907 let new_set = FnvHashSet::from_iter(vec![
908 transformed.to_owned(),
909 {
912 match self {
913 Transform::Day => Datum::date(v + 1),
914 _ => Datum::int(v + 1),
915 }
916 },
917 ]);
918 return Some(AdjustedProjection::Set(new_set));
919 }
920 }
921 _ => {
922 return None;
923 }
924 }
925 };
926 None
927 }
928
929 #[inline]
932 fn try_increment_number(datum: &Datum) -> Result<Datum> {
933 match (datum.data_type(), datum.literal()) {
934 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Ok(Datum::int(v + 1)),
935 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Ok(Datum::long(v + 1)),
936 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
937 Datum::decimal(decimal_from_i128_with_scale(v + 1, 0))
938 }
939 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Ok(Datum::date(v + 1)),
940 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
941 Ok(Datum::timestamp_micros(v + 1))
942 }
943 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
944 Ok(Datum::timestamp_nanos(v + 1))
945 }
946 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
947 Ok(Datum::timestamptz_micros(v + 1))
948 }
949 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
950 Ok(Datum::timestamptz_nanos(v + 1))
951 }
952 (PrimitiveType::Int, _)
953 | (PrimitiveType::Long, _)
954 | (PrimitiveType::Decimal { .. }, _)
955 | (PrimitiveType::Date, _)
956 | (PrimitiveType::Timestamp, _) => Err(Error::new(
957 ErrorKind::Unexpected,
958 format!(
959 "Unsupported literal increment for type: {:?}",
960 datum.data_type()
961 ),
962 )),
963 _ => Ok(datum.to_owned()),
964 }
965 }
966
967 #[inline]
970 fn try_decrement_number(datum: &Datum) -> Result<Datum> {
971 match (datum.data_type(), datum.literal()) {
972 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Ok(Datum::int(v - 1)),
973 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Ok(Datum::long(v - 1)),
974 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
975 Datum::decimal(decimal_from_i128_with_scale(v - 1, 0))
976 }
977 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Ok(Datum::date(v - 1)),
978 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
979 Ok(Datum::timestamp_micros(v - 1))
980 }
981 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
982 Ok(Datum::timestamp_nanos(v - 1))
983 }
984 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
985 Ok(Datum::timestamptz_micros(v - 1))
986 }
987 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
988 Ok(Datum::timestamptz_nanos(v - 1))
989 }
990 (PrimitiveType::Int, _)
991 | (PrimitiveType::Long, _)
992 | (PrimitiveType::Decimal { .. }, _)
993 | (PrimitiveType::Date, _)
994 | (PrimitiveType::Timestamp, _) => Err(Error::new(
995 ErrorKind::Unexpected,
996 format!(
997 "Unsupported literal decrement for type: {:?}",
998 datum.data_type()
999 ),
1000 )),
1001 _ => Ok(datum.to_owned()),
1002 }
1003 }
1004
1005 fn truncate_number_strict(
1006 &self,
1007 name: &str,
1008 expr: &BinaryExpression<BoundReference>,
1009 func: &BoxedTransformFunction,
1010 ) -> Result<Option<Predicate>> {
1011 let boundary = expr.literal();
1012
1013 if !matches!(
1014 boundary.data_type(),
1015 &PrimitiveType::Int
1016 | &PrimitiveType::Long
1017 | &PrimitiveType::Decimal { .. }
1018 | &PrimitiveType::Date
1019 | &PrimitiveType::Timestamp
1020 | &PrimitiveType::Timestamptz
1021 | &PrimitiveType::TimestampNs
1022 | &PrimitiveType::TimestamptzNs
1023 ) {
1024 return Err(invalid_data!(
1025 "Expected a numeric literal, got: {boundary:?}"
1026 ));
1027 }
1028
1029 let predicate = match expr.op() {
1030 PredicateOperator::LessThan => Some(Predicate::Binary(BinaryExpression::new(
1031 PredicateOperator::LessThan,
1032 Reference::new(name),
1033 func.transform_literal_result(boundary)?,
1034 ))),
1035 PredicateOperator::LessThanOrEq => Some(Predicate::Binary(BinaryExpression::new(
1036 PredicateOperator::LessThan,
1037 Reference::new(name),
1038 func.transform_literal_result(&Self::try_increment_number(boundary)?)?,
1039 ))),
1040 PredicateOperator::GreaterThan => Some(Predicate::Binary(BinaryExpression::new(
1041 PredicateOperator::GreaterThan,
1042 Reference::new(name),
1043 func.transform_literal_result(boundary)?,
1044 ))),
1045 PredicateOperator::GreaterThanOrEq => Some(Predicate::Binary(BinaryExpression::new(
1046 PredicateOperator::GreaterThan,
1047 Reference::new(name),
1048 func.transform_literal_result(&Self::try_decrement_number(boundary)?)?,
1049 ))),
1050 PredicateOperator::NotEq => Some(Predicate::Binary(BinaryExpression::new(
1051 PredicateOperator::NotEq,
1052 Reference::new(name),
1053 func.transform_literal_result(boundary)?,
1054 ))),
1055 _ => None,
1056 };
1057
1058 Ok(predicate)
1059 }
1060
1061 fn truncate_array_strict(
1062 &self,
1063 name: &str,
1064 expr: &BinaryExpression<BoundReference>,
1065 func: &BoxedTransformFunction,
1066 ) -> Result<Option<Predicate>> {
1067 let boundary = expr.literal();
1068
1069 match expr.op() {
1070 PredicateOperator::LessThan | PredicateOperator::LessThanOrEq => {
1071 Ok(Some(Predicate::Binary(BinaryExpression::new(
1072 PredicateOperator::LessThan,
1073 Reference::new(name),
1074 func.transform_literal_result(boundary)?,
1075 ))))
1076 }
1077 PredicateOperator::GreaterThan | PredicateOperator::GreaterThanOrEq => {
1078 Ok(Some(Predicate::Binary(BinaryExpression::new(
1079 PredicateOperator::GreaterThan,
1080 Reference::new(name),
1081 func.transform_literal_result(boundary)?,
1082 ))))
1083 }
1084 PredicateOperator::NotEq => Ok(Some(Predicate::Binary(BinaryExpression::new(
1085 PredicateOperator::NotEq,
1086 Reference::new(name),
1087 func.transform_literal_result(boundary)?,
1088 )))),
1089 _ => Ok(None),
1090 }
1091 }
1092}
1093
1094impl Display for Transform {
1095 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1096 match self {
1097 Transform::Identity => write!(f, "identity"),
1098 Transform::Year => write!(f, "year"),
1099 Transform::Month => write!(f, "month"),
1100 Transform::Day => write!(f, "day"),
1101 Transform::Hour => write!(f, "hour"),
1102 Transform::Void => write!(f, "void"),
1103 Transform::Bucket(length) => write!(f, "bucket[{length}]"),
1104 Transform::Truncate(width) => write!(f, "truncate[{width}]"),
1105 Transform::Unknown => write!(f, "unknown"),
1106 }
1107 }
1108}
1109
1110impl FromStr for Transform {
1111 type Err = Error;
1112
1113 fn from_str(s: &str) -> Result<Self> {
1114 let t = match s {
1115 "identity" => Transform::Identity,
1116 "year" => Transform::Year,
1117 "month" => Transform::Month,
1118 "day" => Transform::Day,
1119 "hour" => Transform::Hour,
1120 "void" => Transform::Void,
1121 "unknown" => Transform::Unknown,
1122 v if v.starts_with("bucket") => {
1123 let length = v
1124 .strip_prefix("bucket")
1125 .expect("transform must starts with `bucket`")
1126 .trim_start_matches('[')
1127 .trim_end_matches(']')
1128 .parse()
1129 .map_err(|err| {
1130 invalid_data!("transform bucket type {v:?} is invalid").with_source(err)
1131 })?;
1132
1133 Transform::Bucket(length)
1134 }
1135 v if v.starts_with("truncate") => {
1136 let width = v
1137 .strip_prefix("truncate")
1138 .expect("transform must starts with `truncate`")
1139 .trim_start_matches('[')
1140 .trim_end_matches(']')
1141 .parse()
1142 .map_err(|err| {
1143 invalid_data!("transform truncate type {v:?} is invalid").with_source(err)
1144 })?;
1145
1146 Transform::Truncate(width)
1147 }
1148 v => {
1149 return Err(invalid_data!("transform {v:?} is invalid"));
1150 }
1151 };
1152
1153 Ok(t)
1154 }
1155}
1156
1157impl Serialize for Transform {
1158 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
1159 where S: Serializer {
1160 serializer.serialize_str(format!("{self}").as_str())
1161 }
1162}
1163
1164impl<'de> Deserialize<'de> for Transform {
1165 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
1166 where D: Deserializer<'de> {
1167 let s = String::deserialize(deserializer)?;
1168 s.parse().map_err(<D::Error as serde::de::Error>::custom)
1169 }
1170}
1171
1172#[derive(Debug)]
1175enum AdjustedProjection {
1176 Single(Datum),
1177 Set(FnvHashSet<Datum>),
1178}
1179
1180#[cfg(test)]
1181mod tests {
1182 use super::*;
1183
1184 fn check_boundary(op: PredicateOperator, input: Datum, expected: Datum) {
1185 let result = Transform::adjust_boundary(&op, &input).unwrap().unwrap();
1186 assert_eq!(result, expected);
1187 }
1188
1189 #[test]
1190 fn test_adjust_boundary_timestamp_types() {
1191 for (datum, dec, inc) in [
1192 (
1193 Datum::timestamptz_micros(1000),
1194 Datum::timestamptz_micros(999),
1195 Datum::timestamptz_micros(1001),
1196 ),
1197 (
1198 Datum::timestamp_nanos(5000),
1199 Datum::timestamp_nanos(4999),
1200 Datum::timestamp_nanos(5001),
1201 ),
1202 (
1203 Datum::timestamptz_nanos(5000),
1204 Datum::timestamptz_nanos(4999),
1205 Datum::timestamptz_nanos(5001),
1206 ),
1207 ] {
1208 check_boundary(PredicateOperator::LessThan, datum.clone(), dec);
1209 check_boundary(PredicateOperator::GreaterThan, datum.clone(), inc);
1210 check_boundary(
1211 PredicateOperator::LessThanOrEq,
1212 datum.clone(),
1213 datum.clone(),
1214 );
1215 check_boundary(PredicateOperator::GreaterThanOrEq, datum.clone(), datum);
1216 }
1217 }
1218
1219 fn human(transform: Transform, primitive: PrimitiveType, ordinal: i32) -> String {
1221 transform.to_human_string(&Type::Primitive(primitive), Some(&Literal::int(ordinal)))
1222 }
1223
1224 fn human_int(transform: Transform, ordinal: i32) -> String {
1226 human(transform, PrimitiveType::Int, ordinal)
1227 }
1228
1229 #[test]
1230 fn test_to_human_string_year() {
1231 assert_eq!(human_int(Transform::Year, -1970), "0000");
1232 assert_eq!(human_int(Transform::Year, -1), "1969");
1233 assert_eq!(human_int(Transform::Year, 0), "1970");
1234 assert_eq!(human_int(Transform::Year, 47), "2017");
1235 }
1236
1237 #[test]
1238 fn test_to_human_string_month() {
1239 assert_eq!(human_int(Transform::Month, -1970 * 12), "0000-01");
1240 assert_eq!(human_int(Transform::Month, -13), "1968-12");
1241 assert_eq!(human_int(Transform::Month, -12), "1969-01");
1242 assert_eq!(human_int(Transform::Month, -1), "1969-12");
1243 assert_eq!(human_int(Transform::Month, 0), "1970-01");
1244 assert_eq!(human_int(Transform::Month, 11), "1970-12");
1245 assert_eq!(human_int(Transform::Month, 12), "1971-01");
1246 assert_eq!(human_int(Transform::Month, 569), "2017-06");
1247 }
1248
1249 #[test]
1250 fn test_to_human_string_day() {
1251 assert_eq!(human_int(Transform::Day, -1), "1969-12-31");
1252 assert_eq!(human_int(Transform::Day, 0), "1970-01-01");
1253 assert_eq!(human_int(Transform::Day, 31), "1970-02-01");
1254 assert_eq!(human_int(Transform::Day, 17332), "2017-06-15");
1255 }
1256
1257 #[test]
1258 fn test_to_human_string_hour() {
1259 assert_eq!(human_int(Transform::Hour, -24), "1969-12-31-00");
1260 assert_eq!(human_int(Transform::Hour, -1), "1969-12-31-23");
1261 assert_eq!(human_int(Transform::Hour, 0), "1970-01-01-00");
1262 assert_eq!(human_int(Transform::Hour, 23), "1970-01-01-23");
1263 assert_eq!(human_int(Transform::Hour, 24), "1970-01-02-00");
1264 assert_eq!(human_int(Transform::Hour, 1000), "1970-02-11-16");
1265 assert_eq!(human_int(Transform::Hour, 415984), "2017-06-15-16");
1266 }
1267
1268 #[test]
1272 fn test_to_human_string_ignores_declared_type_for_temporal_transforms() {
1273 assert_eq!(human(Transform::Year, PrimitiveType::Date, 47), "2017");
1274 assert_eq!(human(Transform::Month, PrimitiveType::Date, 569), "2017-06");
1275 assert_eq!(
1276 human(Transform::Day, PrimitiveType::Int, 17332),
1277 "2017-06-15"
1278 );
1279 assert_eq!(
1280 human(Transform::Hour, PrimitiveType::Date, 415984),
1281 "2017-06-15-16"
1282 );
1283 }
1284
1285 #[test]
1288 fn test_to_human_string_defers_to_datum_for_other_transforms() {
1289 assert_eq!(
1290 human(Transform::Identity, PrimitiveType::Int, 17332),
1291 "17332"
1292 );
1293 assert_eq!(
1294 human(Transform::Identity, PrimitiveType::Date, 17332),
1295 "2017-06-15"
1296 );
1297 assert_eq!(human(Transform::Bucket(16), PrimitiveType::Int, 5), "5");
1298
1299 assert_eq!(
1302 Transform::Year.to_human_string(
1303 &Type::Primitive(PrimitiveType::String),
1304 Some(&Literal::string("unformatted"))
1305 ),
1306 "unformatted"
1307 );
1308 }
1309
1310 #[test]
1311 fn test_to_human_string_null_cases() {
1312 assert_eq!(human(Transform::Void, PrimitiveType::Int, 47), "null");
1313 assert_eq!(human(Transform::Void, PrimitiveType::Date, 17332), "null");
1314 for transform in [
1315 Transform::Year,
1316 Transform::Month,
1317 Transform::Day,
1318 Transform::Hour,
1319 Transform::Identity,
1320 Transform::Void,
1321 ] {
1322 assert_eq!(
1323 transform.to_human_string(&Type::Primitive(PrimitiveType::Int), None),
1324 "null"
1325 );
1326 }
1327 }
1328}