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