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};
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(Error::new(
286 ErrorKind::DataInvalid,
287 format!("{input_type} is not a valid input type of identity transform",),
288 ))
289 }
290 }
291 Transform::Void => Ok(input_type.clone()),
292 Transform::Unknown => Ok(Type::Primitive(PrimitiveType::String)),
293 Transform::Bucket(_) => {
294 if let Type::Primitive(p) = input_type {
295 match p {
296 PrimitiveType::Int
297 | PrimitiveType::Long
298 | PrimitiveType::Decimal { .. }
299 | PrimitiveType::Date
300 | PrimitiveType::Time
301 | PrimitiveType::Timestamp
302 | PrimitiveType::Timestamptz
303 | PrimitiveType::TimestampNs
304 | PrimitiveType::TimestamptzNs
305 | PrimitiveType::String
306 | PrimitiveType::Uuid
307 | PrimitiveType::Fixed(_)
308 | PrimitiveType::Binary => Ok(Type::Primitive(PrimitiveType::Int)),
309 _ => Err(Error::new(
310 ErrorKind::DataInvalid,
311 format!("{input_type} is not a valid input type of bucket transform",),
312 )),
313 }
314 } else {
315 Err(Error::new(
316 ErrorKind::DataInvalid,
317 format!("{input_type} is not a valid input type of bucket transform",),
318 ))
319 }
320 }
321 Transform::Truncate(_) => {
322 if let Type::Primitive(p) = input_type {
323 match p {
324 PrimitiveType::Int
325 | PrimitiveType::Long
326 | PrimitiveType::String
327 | PrimitiveType::Binary
328 | PrimitiveType::Decimal { .. } => Ok(input_type.clone()),
329 _ => Err(Error::new(
330 ErrorKind::DataInvalid,
331 format!("{input_type} is not a valid input type of truncate transform",),
332 )),
333 }
334 } else {
335 Err(Error::new(
336 ErrorKind::DataInvalid,
337 format!("{input_type} is not a valid input type of truncate transform",),
338 ))
339 }
340 }
341 Transform::Year | Transform::Month => {
342 if let Type::Primitive(p) = input_type {
343 match p {
344 PrimitiveType::Timestamp
345 | PrimitiveType::Timestamptz
346 | PrimitiveType::TimestampNs
347 | PrimitiveType::TimestamptzNs
348 | PrimitiveType::Date => Ok(Type::Primitive(PrimitiveType::Int)),
349 _ => Err(Error::new(
350 ErrorKind::DataInvalid,
351 format!("{input_type} is not a valid input type of {self} transform",),
352 )),
353 }
354 } else {
355 Err(Error::new(
356 ErrorKind::DataInvalid,
357 format!("{input_type} is not a valid input type of {self} transform",),
358 ))
359 }
360 }
361 Transform::Day => {
362 if let Type::Primitive(p) = input_type {
363 match p {
364 PrimitiveType::Timestamp
365 | PrimitiveType::Timestamptz
366 | PrimitiveType::TimestampNs
367 | PrimitiveType::TimestamptzNs
368 | PrimitiveType::Date => Ok(Type::Primitive(PrimitiveType::Date)),
369 _ => Err(Error::new(
370 ErrorKind::DataInvalid,
371 format!("{input_type} is not a valid input type of {self} transform",),
372 )),
373 }
374 } else {
375 Err(Error::new(
376 ErrorKind::DataInvalid,
377 format!("{input_type} is not a valid input type of {self} transform",),
378 ))
379 }
380 }
381 Transform::Hour => {
382 if let Type::Primitive(p) = input_type {
383 match p {
384 PrimitiveType::Timestamp
385 | PrimitiveType::Timestamptz
386 | PrimitiveType::TimestampNs
387 | PrimitiveType::TimestamptzNs => Ok(Type::Primitive(PrimitiveType::Int)),
388 _ => Err(Error::new(
389 ErrorKind::DataInvalid,
390 format!("{input_type} is not a valid input type of {self} transform",),
391 )),
392 }
393 } else {
394 Err(Error::new(
395 ErrorKind::DataInvalid,
396 format!("{input_type} is not a valid input type of {self} transform",),
397 ))
398 }
399 }
400 }
401 }
402
403 pub fn preserves_order(&self) -> bool {
405 !matches!(
406 self,
407 Transform::Void | Transform::Bucket(_) | Transform::Unknown
408 )
409 }
410
411 pub fn dedup_name(&self) -> String {
414 match self {
415 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
416 "time".to_string()
417 }
418 _ => format!("{self}"),
419 }
420 }
421
422 pub fn satisfies_order_of(&self, other: &Self) -> bool {
428 match self {
429 Transform::Identity => other.preserves_order(),
430 Transform::Hour => matches!(
431 other,
432 Transform::Hour | Transform::Day | Transform::Month | Transform::Year
433 ),
434 Transform::Day => matches!(other, Transform::Day | Transform::Month | Transform::Year),
435 Transform::Month => matches!(other, Transform::Month | Transform::Year),
436 _ => self == other,
437 }
438 }
439
440 pub fn strict_project(
453 &self,
454 name: &str,
455 predicate: &BoundPredicate,
456 ) -> Result<Option<Predicate>> {
457 let func = create_transform_function(self)?;
458
459 match self {
460 Transform::Identity => match predicate {
461 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
462 BoundPredicate::Binary(expr) => Ok(Some(Predicate::Binary(BinaryExpression::new(
463 expr.op(),
464 Reference::new(name),
465 expr.literal().to_owned(),
466 )))),
467 BoundPredicate::Set(expr) => Ok(Some(Predicate::Set(SetExpression::new(
468 expr.op(),
469 Reference::new(name),
470 expr.literals().to_owned(),
471 )))),
472 _ => Ok(None),
473 },
474 Transform::Bucket(_) => match predicate {
475 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
476 BoundPredicate::Binary(expr) => {
477 self.project_binary_expr(name, PredicateOperator::NotEq, expr, &func)
478 }
479 BoundPredicate::Set(expr) => {
480 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
481 }
482 _ => Ok(None),
483 },
484 Transform::Truncate(width) => match predicate {
485 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
486 BoundPredicate::Binary(expr) => {
487 if matches!(
488 expr.term().field().field_type.as_primitive_type(),
489 Some(&PrimitiveType::Int)
490 | Some(&PrimitiveType::Long)
491 | Some(&PrimitiveType::Decimal { .. })
492 ) {
493 self.truncate_number_strict(name, expr, &func)
494 } else if expr.op() == PredicateOperator::StartsWith {
495 let len = match expr.literal().literal() {
496 PrimitiveLiteral::String(s) => s.len(),
497 PrimitiveLiteral::Binary(b) => b.len(),
498 _ => {
499 return Err(Error::new(
500 ErrorKind::DataInvalid,
501 format!(
502 "Expected a string or binary literal, got: {:?}",
503 expr.literal()
504 ),
505 ));
506 }
507 };
508 match len.cmp(&(*width as usize)) {
509 Ordering::Less => Ok(Some(Predicate::Binary(BinaryExpression::new(
510 PredicateOperator::StartsWith,
511 Reference::new(name),
512 expr.literal().to_owned(),
513 )))),
514 Ordering::Equal => Ok(Some(Predicate::Binary(BinaryExpression::new(
515 PredicateOperator::Eq,
516 Reference::new(name),
517 expr.literal().to_owned(),
518 )))),
519 Ordering::Greater => Ok(None),
520 }
521 } else if expr.op() == PredicateOperator::NotStartsWith {
522 let len = match expr.literal().literal() {
523 PrimitiveLiteral::String(s) => s.len(),
524 PrimitiveLiteral::Binary(b) => b.len(),
525 _ => {
526 return Err(Error::new(
527 ErrorKind::DataInvalid,
528 format!(
529 "Expected a string or binary literal, got: {:?}",
530 expr.literal()
531 ),
532 ));
533 }
534 };
535 match len.cmp(&(*width as usize)) {
536 Ordering::Less => Ok(Some(Predicate::Binary(BinaryExpression::new(
537 PredicateOperator::NotStartsWith,
538 Reference::new(name),
539 expr.literal().to_owned(),
540 )))),
541 Ordering::Equal => Ok(Some(Predicate::Binary(BinaryExpression::new(
542 PredicateOperator::NotEq,
543 Reference::new(name),
544 expr.literal().to_owned(),
545 )))),
546 Ordering::Greater => {
547 Ok(Some(Predicate::Binary(BinaryExpression::new(
548 expr.op(),
549 Reference::new(name),
550 func.transform_literal_result(expr.literal())?,
551 ))))
552 }
553 }
554 } else {
555 self.truncate_array_strict(name, expr, &func)
556 }
557 }
558 BoundPredicate::Set(expr) => {
559 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
560 }
561 _ => Ok(None),
562 },
563 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
564 match predicate {
565 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
566 BoundPredicate::Binary(expr) => self.truncate_number_strict(name, expr, &func),
567 BoundPredicate::Set(expr) => {
568 self.project_set_expr(expr, PredicateOperator::NotIn, name, &func)
569 }
570 _ => Ok(None),
571 }
572 }
573 _ => Ok(None),
574 }
575 }
576
577 pub fn project(&self, name: &str, predicate: &BoundPredicate) -> Result<Option<Predicate>> {
590 let func = create_transform_function(self)?;
591
592 match self {
593 Transform::Identity => match predicate {
594 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
595 BoundPredicate::Binary(expr) => Ok(Some(Predicate::Binary(BinaryExpression::new(
596 expr.op(),
597 Reference::new(name),
598 expr.literal().to_owned(),
599 )))),
600 BoundPredicate::Set(expr) => Ok(Some(Predicate::Set(SetExpression::new(
601 expr.op(),
602 Reference::new(name),
603 expr.literals().to_owned(),
604 )))),
605 _ => Ok(None),
606 },
607 Transform::Bucket(_) => match predicate {
608 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
609 BoundPredicate::Binary(expr) => {
610 self.project_binary_expr(name, PredicateOperator::Eq, expr, &func)
611 }
612 BoundPredicate::Set(expr) => {
613 self.project_set_expr(expr, PredicateOperator::In, name, &func)
614 }
615 _ => Ok(None),
616 },
617 Transform::Truncate(width) => match predicate {
618 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
619 BoundPredicate::Binary(expr) => {
620 self.project_binary_with_adjusted_boundary(name, expr, &func, Some(*width))
621 }
622 BoundPredicate::Set(expr) => {
623 self.project_set_expr(expr, PredicateOperator::In, name, &func)
624 }
625 _ => Ok(None),
626 },
627 Transform::Year | Transform::Month | Transform::Day | Transform::Hour => {
628 match predicate {
629 BoundPredicate::Unary(expr) => Self::project_unary(expr.op(), name),
630 BoundPredicate::Binary(expr) => {
631 self.project_binary_with_adjusted_boundary(name, expr, &func, None)
632 }
633 BoundPredicate::Set(expr) => {
634 self.project_set_expr(expr, PredicateOperator::In, name, &func)
635 }
636 _ => Ok(None),
637 }
638 }
639 _ => Ok(None),
640 }
641 }
642
643 fn can_transform(&self, datum: &Datum) -> bool {
645 let input_type = datum.data_type().clone();
646 self.result_type(&Type::Primitive(input_type)).is_ok()
647 }
648
649 fn project_unary(op: PredicateOperator, name: &str) -> Result<Option<Predicate>> {
651 Ok(Some(Predicate::Unary(UnaryExpression::new(
652 op,
653 Reference::new(name),
654 ))))
655 }
656
657 fn project_binary_expr(
664 &self,
665 name: &str,
666 op: PredicateOperator,
667 expr: &BinaryExpression<BoundReference>,
668 func: &BoxedTransformFunction,
669 ) -> Result<Option<Predicate>> {
670 if expr.op() != op || !self.can_transform(expr.literal()) {
671 return Ok(None);
672 }
673
674 Ok(Some(Predicate::Binary(BinaryExpression::new(
675 expr.op(),
676 Reference::new(name),
677 func.transform_literal_result(expr.literal())?,
678 ))))
679 }
680
681 fn project_binary_with_adjusted_boundary(
690 &self,
691 name: &str,
692 expr: &BinaryExpression<BoundReference>,
693 func: &BoxedTransformFunction,
694 width: Option<u32>,
695 ) -> Result<Option<Predicate>> {
696 if !self.can_transform(expr.literal()) {
697 return Ok(None);
698 }
699
700 let op = &expr.op();
701 let datum = &expr.literal();
702
703 if let Some(boundary) = Self::adjust_boundary(op, datum)? {
704 let transformed_projection = func.transform_literal_result(&boundary)?;
705
706 let adjusted_projection =
707 self.adjust_time_projection(op, datum, &transformed_projection);
708
709 let adjusted_operator = Self::adjust_operator(op, datum, width);
710
711 if let Some(op) = adjusted_operator {
712 let predicate = match adjusted_projection {
713 None => Predicate::Binary(BinaryExpression::new(
714 op,
715 Reference::new(name),
716 transformed_projection,
717 )),
718 Some(AdjustedProjection::Single(d)) => {
719 Predicate::Binary(BinaryExpression::new(op, Reference::new(name), d))
720 }
721 Some(AdjustedProjection::Set(d)) => Predicate::Set(SetExpression::new(
722 PredicateOperator::In,
723 Reference::new(name),
724 d,
725 )),
726 };
727 return Ok(Some(predicate));
728 }
729 };
730
731 Ok(None)
732 }
733
734 fn project_set_expr(
737 &self,
738 expr: &SetExpression<BoundReference>,
739 op: PredicateOperator,
740 name: &str,
741 func: &BoxedTransformFunction,
742 ) -> Result<Option<Predicate>> {
743 if expr.op() != op || expr.literals().iter().any(|d| !self.can_transform(d)) {
744 return Ok(None);
745 }
746
747 let mut new_set = FnvHashSet::default();
748
749 for lit in expr.literals() {
750 let datum = func.transform_literal_result(lit)?;
751
752 if let Some(AdjustedProjection::Single(d)) =
753 self.adjust_time_projection(&op, lit, &datum)
754 {
755 new_set.insert(d);
756 };
757
758 new_set.insert(datum);
759 }
760
761 Ok(Some(Predicate::Set(SetExpression::new(
762 expr.op(),
763 Reference::new(name),
764 new_set,
765 ))))
766 }
767
768 fn adjust_boundary(op: &PredicateOperator, datum: &Datum) -> Result<Option<Datum>> {
778 let adjusted_boundary = match op {
779 PredicateOperator::LessThan => match (datum.data_type(), datum.literal()) {
780 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Some(Datum::int(v - 1)),
781 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Some(Datum::long(v - 1)),
782 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
783 Some(Datum::decimal(decimal_from_i128_with_scale(v - 1, 0))?)
784 }
785 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Some(Datum::date(v - 1)),
786 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
787 Some(Datum::timestamp_micros(v - 1))
788 }
789 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
790 Some(Datum::timestamptz_micros(v - 1))
791 }
792 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
793 Some(Datum::timestamp_nanos(v - 1))
794 }
795 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
796 Some(Datum::timestamptz_nanos(v - 1))
797 }
798 _ => Some(datum.to_owned()),
799 },
800 PredicateOperator::GreaterThan => match (datum.data_type(), datum.literal()) {
801 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Some(Datum::int(v + 1)),
802 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Some(Datum::long(v + 1)),
803 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
804 Some(Datum::decimal(decimal_from_i128_with_scale(v + 1, 0))?)
805 }
806 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Some(Datum::date(v + 1)),
807 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
808 Some(Datum::timestamp_micros(v + 1))
809 }
810 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
811 Some(Datum::timestamptz_micros(v + 1))
812 }
813 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
814 Some(Datum::timestamp_nanos(v + 1))
815 }
816 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
817 Some(Datum::timestamptz_nanos(v + 1))
818 }
819 _ => Some(datum.to_owned()),
820 },
821 PredicateOperator::Eq
822 | PredicateOperator::LessThanOrEq
823 | PredicateOperator::GreaterThanOrEq
824 | PredicateOperator::StartsWith
825 | PredicateOperator::NotStartsWith => Some(datum.to_owned()),
826 _ => None,
827 };
828
829 Ok(adjusted_boundary)
830 }
831
832 fn adjust_operator(
843 op: &PredicateOperator,
844 datum: &Datum,
845 width: Option<u32>,
846 ) -> Option<PredicateOperator> {
847 match op {
848 PredicateOperator::LessThan => Some(PredicateOperator::LessThanOrEq),
849 PredicateOperator::GreaterThan => Some(PredicateOperator::GreaterThanOrEq),
850 PredicateOperator::StartsWith => match datum.literal() {
851 PrimitiveLiteral::String(s) => {
852 if let Some(w) = width
853 && s.len() == w as usize
854 {
855 return Some(PredicateOperator::Eq);
856 };
857 Some(*op)
858 }
859 _ => Some(*op),
860 },
861 PredicateOperator::NotStartsWith => match datum.literal() {
862 PrimitiveLiteral::String(s) => {
863 if let Some(w) = width {
864 let w = w as usize;
865
866 if s.len() == w {
867 return Some(PredicateOperator::NotEq);
868 }
869
870 if s.len() < w {
871 return Some(*op);
872 }
873
874 return None;
875 };
876 Some(*op)
877 }
878 _ => Some(*op),
879 },
880 _ => Some(*op),
881 }
882 }
883
884 fn adjust_time_projection(
887 &self,
888 op: &PredicateOperator,
889 original: &Datum,
890 transformed: &Datum,
891 ) -> Option<AdjustedProjection> {
892 let should_adjust = match self {
893 Transform::Day => matches!(original.data_type(), PrimitiveType::Timestamp),
894 Transform::Year | Transform::Month => true,
895 _ => false,
896 };
897
898 if should_adjust && let &PrimitiveLiteral::Int(v) = transformed.literal() {
899 match op {
900 PredicateOperator::LessThan
901 | PredicateOperator::LessThanOrEq
902 | PredicateOperator::In => {
903 if v < 0 {
904 match self {
907 Transform::Day => {
908 return Some(AdjustedProjection::Single(Datum::date(v + 1)));
909 }
910 _ => {
911 return Some(AdjustedProjection::Single(Datum::int(v + 1)));
912 }
913 }
914 };
915 }
916 PredicateOperator::Eq => {
917 if v < 0 {
918 let new_set = FnvHashSet::from_iter(vec![
919 transformed.to_owned(),
920 {
923 match self {
924 Transform::Day => Datum::date(v + 1),
925 _ => Datum::int(v + 1),
926 }
927 },
928 ]);
929 return Some(AdjustedProjection::Set(new_set));
930 }
931 }
932 _ => {
933 return None;
934 }
935 }
936 };
937 None
938 }
939
940 #[inline]
943 fn try_increment_number(datum: &Datum) -> Result<Datum> {
944 match (datum.data_type(), datum.literal()) {
945 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Ok(Datum::int(v + 1)),
946 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Ok(Datum::long(v + 1)),
947 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
948 Datum::decimal(decimal_from_i128_with_scale(v + 1, 0))
949 }
950 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Ok(Datum::date(v + 1)),
951 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
952 Ok(Datum::timestamp_micros(v + 1))
953 }
954 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
955 Ok(Datum::timestamp_nanos(v + 1))
956 }
957 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
958 Ok(Datum::timestamptz_micros(v + 1))
959 }
960 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
961 Ok(Datum::timestamptz_nanos(v + 1))
962 }
963 (PrimitiveType::Int, _)
964 | (PrimitiveType::Long, _)
965 | (PrimitiveType::Decimal { .. }, _)
966 | (PrimitiveType::Date, _)
967 | (PrimitiveType::Timestamp, _) => Err(Error::new(
968 ErrorKind::Unexpected,
969 format!(
970 "Unsupported literal increment for type: {:?}",
971 datum.data_type()
972 ),
973 )),
974 _ => Ok(datum.to_owned()),
975 }
976 }
977
978 #[inline]
981 fn try_decrement_number(datum: &Datum) -> Result<Datum> {
982 match (datum.data_type(), datum.literal()) {
983 (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Ok(Datum::int(v - 1)),
984 (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Ok(Datum::long(v - 1)),
985 (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => {
986 Datum::decimal(decimal_from_i128_with_scale(v - 1, 0))
987 }
988 (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Ok(Datum::date(v - 1)),
989 (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
990 Ok(Datum::timestamp_micros(v - 1))
991 }
992 (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => {
993 Ok(Datum::timestamp_nanos(v - 1))
994 }
995 (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => {
996 Ok(Datum::timestamptz_micros(v - 1))
997 }
998 (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => {
999 Ok(Datum::timestamptz_nanos(v - 1))
1000 }
1001 (PrimitiveType::Int, _)
1002 | (PrimitiveType::Long, _)
1003 | (PrimitiveType::Decimal { .. }, _)
1004 | (PrimitiveType::Date, _)
1005 | (PrimitiveType::Timestamp, _) => Err(Error::new(
1006 ErrorKind::Unexpected,
1007 format!(
1008 "Unsupported literal decrement for type: {:?}",
1009 datum.data_type()
1010 ),
1011 )),
1012 _ => Ok(datum.to_owned()),
1013 }
1014 }
1015
1016 fn truncate_number_strict(
1017 &self,
1018 name: &str,
1019 expr: &BinaryExpression<BoundReference>,
1020 func: &BoxedTransformFunction,
1021 ) -> Result<Option<Predicate>> {
1022 let boundary = expr.literal();
1023
1024 if !matches!(
1025 boundary.data_type(),
1026 &PrimitiveType::Int
1027 | &PrimitiveType::Long
1028 | &PrimitiveType::Decimal { .. }
1029 | &PrimitiveType::Date
1030 | &PrimitiveType::Timestamp
1031 | &PrimitiveType::Timestamptz
1032 | &PrimitiveType::TimestampNs
1033 | &PrimitiveType::TimestamptzNs
1034 ) {
1035 return Err(Error::new(
1036 ErrorKind::DataInvalid,
1037 format!("Expected a numeric literal, got: {boundary:?}"),
1038 ));
1039 }
1040
1041 let predicate = match expr.op() {
1042 PredicateOperator::LessThan => Some(Predicate::Binary(BinaryExpression::new(
1043 PredicateOperator::LessThan,
1044 Reference::new(name),
1045 func.transform_literal_result(boundary)?,
1046 ))),
1047 PredicateOperator::LessThanOrEq => Some(Predicate::Binary(BinaryExpression::new(
1048 PredicateOperator::LessThan,
1049 Reference::new(name),
1050 func.transform_literal_result(&Self::try_increment_number(boundary)?)?,
1051 ))),
1052 PredicateOperator::GreaterThan => Some(Predicate::Binary(BinaryExpression::new(
1053 PredicateOperator::GreaterThan,
1054 Reference::new(name),
1055 func.transform_literal_result(boundary)?,
1056 ))),
1057 PredicateOperator::GreaterThanOrEq => Some(Predicate::Binary(BinaryExpression::new(
1058 PredicateOperator::GreaterThan,
1059 Reference::new(name),
1060 func.transform_literal_result(&Self::try_decrement_number(boundary)?)?,
1061 ))),
1062 PredicateOperator::NotEq => Some(Predicate::Binary(BinaryExpression::new(
1063 PredicateOperator::NotEq,
1064 Reference::new(name),
1065 func.transform_literal_result(boundary)?,
1066 ))),
1067 _ => None,
1068 };
1069
1070 Ok(predicate)
1071 }
1072
1073 fn truncate_array_strict(
1074 &self,
1075 name: &str,
1076 expr: &BinaryExpression<BoundReference>,
1077 func: &BoxedTransformFunction,
1078 ) -> Result<Option<Predicate>> {
1079 let boundary = expr.literal();
1080
1081 match expr.op() {
1082 PredicateOperator::LessThan | PredicateOperator::LessThanOrEq => {
1083 Ok(Some(Predicate::Binary(BinaryExpression::new(
1084 PredicateOperator::LessThan,
1085 Reference::new(name),
1086 func.transform_literal_result(boundary)?,
1087 ))))
1088 }
1089 PredicateOperator::GreaterThan | PredicateOperator::GreaterThanOrEq => {
1090 Ok(Some(Predicate::Binary(BinaryExpression::new(
1091 PredicateOperator::GreaterThan,
1092 Reference::new(name),
1093 func.transform_literal_result(boundary)?,
1094 ))))
1095 }
1096 PredicateOperator::NotEq => Ok(Some(Predicate::Binary(BinaryExpression::new(
1097 PredicateOperator::NotEq,
1098 Reference::new(name),
1099 func.transform_literal_result(boundary)?,
1100 )))),
1101 _ => Ok(None),
1102 }
1103 }
1104}
1105
1106impl Display for Transform {
1107 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1108 match self {
1109 Transform::Identity => write!(f, "identity"),
1110 Transform::Year => write!(f, "year"),
1111 Transform::Month => write!(f, "month"),
1112 Transform::Day => write!(f, "day"),
1113 Transform::Hour => write!(f, "hour"),
1114 Transform::Void => write!(f, "void"),
1115 Transform::Bucket(length) => write!(f, "bucket[{length}]"),
1116 Transform::Truncate(width) => write!(f, "truncate[{width}]"),
1117 Transform::Unknown => write!(f, "unknown"),
1118 }
1119 }
1120}
1121
1122impl FromStr for Transform {
1123 type Err = Error;
1124
1125 fn from_str(s: &str) -> Result<Self> {
1126 let t = match s {
1127 "identity" => Transform::Identity,
1128 "year" => Transform::Year,
1129 "month" => Transform::Month,
1130 "day" => Transform::Day,
1131 "hour" => Transform::Hour,
1132 "void" => Transform::Void,
1133 "unknown" => Transform::Unknown,
1134 v if v.starts_with("bucket") => {
1135 let length = v
1136 .strip_prefix("bucket")
1137 .expect("transform must starts with `bucket`")
1138 .trim_start_matches('[')
1139 .trim_end_matches(']')
1140 .parse()
1141 .map_err(|err| {
1142 Error::new(
1143 ErrorKind::DataInvalid,
1144 format!("transform bucket type {v:?} is invalid"),
1145 )
1146 .with_source(err)
1147 })?;
1148
1149 Transform::Bucket(length)
1150 }
1151 v if v.starts_with("truncate") => {
1152 let width = v
1153 .strip_prefix("truncate")
1154 .expect("transform must starts with `truncate`")
1155 .trim_start_matches('[')
1156 .trim_end_matches(']')
1157 .parse()
1158 .map_err(|err| {
1159 Error::new(
1160 ErrorKind::DataInvalid,
1161 format!("transform truncate type {v:?} is invalid"),
1162 )
1163 .with_source(err)
1164 })?;
1165
1166 Transform::Truncate(width)
1167 }
1168 v => {
1169 return Err(Error::new(
1170 ErrorKind::DataInvalid,
1171 format!("transform {v:?} is invalid"),
1172 ));
1173 }
1174 };
1175
1176 Ok(t)
1177 }
1178}
1179
1180impl Serialize for Transform {
1181 fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
1182 where S: Serializer {
1183 serializer.serialize_str(format!("{self}").as_str())
1184 }
1185}
1186
1187impl<'de> Deserialize<'de> for Transform {
1188 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
1189 where D: Deserializer<'de> {
1190 let s = String::deserialize(deserializer)?;
1191 s.parse().map_err(<D::Error as serde::de::Error>::custom)
1192 }
1193}
1194
1195#[derive(Debug)]
1198enum AdjustedProjection {
1199 Single(Datum),
1200 Set(FnvHashSet<Datum>),
1201}
1202
1203#[cfg(test)]
1204mod tests {
1205 use super::*;
1206
1207 fn check_boundary(op: PredicateOperator, input: Datum, expected: Datum) {
1208 let result = Transform::adjust_boundary(&op, &input).unwrap().unwrap();
1209 assert_eq!(result, expected);
1210 }
1211
1212 #[test]
1213 fn test_adjust_boundary_timestamp_types() {
1214 for (datum, dec, inc) in [
1215 (
1216 Datum::timestamptz_micros(1000),
1217 Datum::timestamptz_micros(999),
1218 Datum::timestamptz_micros(1001),
1219 ),
1220 (
1221 Datum::timestamp_nanos(5000),
1222 Datum::timestamp_nanos(4999),
1223 Datum::timestamp_nanos(5001),
1224 ),
1225 (
1226 Datum::timestamptz_nanos(5000),
1227 Datum::timestamptz_nanos(4999),
1228 Datum::timestamptz_nanos(5001),
1229 ),
1230 ] {
1231 check_boundary(PredicateOperator::LessThan, datum.clone(), dec);
1232 check_boundary(PredicateOperator::GreaterThan, datum.clone(), inc);
1233 check_boundary(
1234 PredicateOperator::LessThanOrEq,
1235 datum.clone(),
1236 datum.clone(),
1237 );
1238 check_boundary(PredicateOperator::GreaterThanOrEq, datum.clone(), datum);
1239 }
1240 }
1241
1242 fn human(transform: Transform, primitive: PrimitiveType, ordinal: i32) -> String {
1244 transform.to_human_string(&Type::Primitive(primitive), Some(&Literal::int(ordinal)))
1245 }
1246
1247 fn human_int(transform: Transform, ordinal: i32) -> String {
1249 human(transform, PrimitiveType::Int, ordinal)
1250 }
1251
1252 #[test]
1253 fn test_to_human_string_year() {
1254 assert_eq!(human_int(Transform::Year, -1970), "0000");
1255 assert_eq!(human_int(Transform::Year, -1), "1969");
1256 assert_eq!(human_int(Transform::Year, 0), "1970");
1257 assert_eq!(human_int(Transform::Year, 47), "2017");
1258 }
1259
1260 #[test]
1261 fn test_to_human_string_month() {
1262 assert_eq!(human_int(Transform::Month, -1970 * 12), "0000-01");
1263 assert_eq!(human_int(Transform::Month, -13), "1968-12");
1264 assert_eq!(human_int(Transform::Month, -12), "1969-01");
1265 assert_eq!(human_int(Transform::Month, -1), "1969-12");
1266 assert_eq!(human_int(Transform::Month, 0), "1970-01");
1267 assert_eq!(human_int(Transform::Month, 11), "1970-12");
1268 assert_eq!(human_int(Transform::Month, 12), "1971-01");
1269 assert_eq!(human_int(Transform::Month, 569), "2017-06");
1270 }
1271
1272 #[test]
1273 fn test_to_human_string_day() {
1274 assert_eq!(human_int(Transform::Day, -1), "1969-12-31");
1275 assert_eq!(human_int(Transform::Day, 0), "1970-01-01");
1276 assert_eq!(human_int(Transform::Day, 31), "1970-02-01");
1277 assert_eq!(human_int(Transform::Day, 17332), "2017-06-15");
1278 }
1279
1280 #[test]
1281 fn test_to_human_string_hour() {
1282 assert_eq!(human_int(Transform::Hour, -24), "1969-12-31-00");
1283 assert_eq!(human_int(Transform::Hour, -1), "1969-12-31-23");
1284 assert_eq!(human_int(Transform::Hour, 0), "1970-01-01-00");
1285 assert_eq!(human_int(Transform::Hour, 23), "1970-01-01-23");
1286 assert_eq!(human_int(Transform::Hour, 24), "1970-01-02-00");
1287 assert_eq!(human_int(Transform::Hour, 1000), "1970-02-11-16");
1288 assert_eq!(human_int(Transform::Hour, 415984), "2017-06-15-16");
1289 }
1290
1291 #[test]
1295 fn test_to_human_string_ignores_declared_type_for_temporal_transforms() {
1296 assert_eq!(human(Transform::Year, PrimitiveType::Date, 47), "2017");
1297 assert_eq!(human(Transform::Month, PrimitiveType::Date, 569), "2017-06");
1298 assert_eq!(
1299 human(Transform::Day, PrimitiveType::Int, 17332),
1300 "2017-06-15"
1301 );
1302 assert_eq!(
1303 human(Transform::Hour, PrimitiveType::Date, 415984),
1304 "2017-06-15-16"
1305 );
1306 }
1307
1308 #[test]
1311 fn test_to_human_string_defers_to_datum_for_other_transforms() {
1312 assert_eq!(
1313 human(Transform::Identity, PrimitiveType::Int, 17332),
1314 "17332"
1315 );
1316 assert_eq!(
1317 human(Transform::Identity, PrimitiveType::Date, 17332),
1318 "2017-06-15"
1319 );
1320 assert_eq!(human(Transform::Bucket(16), PrimitiveType::Int, 5), "5");
1321
1322 assert_eq!(
1325 Transform::Year.to_human_string(
1326 &Type::Primitive(PrimitiveType::String),
1327 Some(&Literal::string("unformatted"))
1328 ),
1329 "unformatted"
1330 );
1331 }
1332
1333 #[test]
1334 fn test_to_human_string_null_cases() {
1335 assert_eq!(human(Transform::Void, PrimitiveType::Int, 47), "null");
1336 assert_eq!(human(Transform::Void, PrimitiveType::Date, 17332), "null");
1337 for transform in [
1338 Transform::Year,
1339 Transform::Month,
1340 Transform::Day,
1341 Transform::Hour,
1342 Transform::Identity,
1343 Transform::Void,
1344 ] {
1345 assert_eq!(
1346 transform.to_human_string(&Type::Primitive(PrimitiveType::Int), None),
1347 "null"
1348 );
1349 }
1350 }
1351}