1use crate::exact::{calendar_without_out_of_range, stored_out_of_range};
19use polars::chunked_array::cast::CastOptions;
20use polars::prelude::*;
21
22pub fn can_leave_calendar(dtype: &DataType) -> bool {
25 matches!(
26 dtype,
27 DataType::Date | DataType::Datetime(TimeUnit::Milliseconds | TimeUnit::Microseconds, _)
28 )
29}
30
31pub fn text_or_stored(
34 series: &Series,
35 text: impl Fn(&Series) -> PolarsResult<StringChunked>,
36) -> PolarsResult<Series> {
37 let dtype = series.dtype();
38 let text = match calendar_without_out_of_range(series)? {
39 Some(shown) => {
41 let stored = series.to_physical_repr().cast(&DataType::Int64)?;
42 text(&shown)?
43 .iter()
44 .zip(stored.i64()?.iter())
45 .map(|(text, v)| match text {
46 Some(text) => Some(text.to_string()),
47 None => v.and_then(|v| stored_out_of_range(dtype, v)),
48 })
49 .collect::<StringChunked>()
50 }
51 None => text(series)?,
52 };
53 Ok(text.with_name(series.name().clone()).into_series())
54}
55
56pub fn cast_text(series: &Series, options: CastOptions) -> PolarsResult<Series> {
59 if !can_leave_calendar(series.dtype()) {
60 return series.cast_with_options(&DataType::String, options);
61 }
62 text_or_stored(series, |s| Ok(s.cast(&DataType::String)?.str()?.clone()))
63}
64
65pub fn formatted(series: &Series, format: &str) -> PolarsResult<Series> {
68 if !can_leave_calendar(series.dtype()) {
69 return TemporalMethods::to_string(series, format);
70 }
71 text_or_stored(series, |s| {
72 Ok(TemporalMethods::to_string(s, format)?.str()?.clone())
73 })
74}
75
76fn text_field(_: &Schema, field: &Field) -> PolarsResult<Field> {
77 Ok(Field::new(field.name().clone(), DataType::String))
78}
79
80pub fn text_expr(expr: Expr, options: CastOptions) -> Expr {
82 expr.map_with_fmt_str(
83 move |c| cast_text(c.as_materialized_series(), options).map(Column::from),
84 text_field,
85 "past_calendar_text",
86 )
87}
88
89pub fn format_expr(expr: Expr, format: String) -> Expr {
91 expr.map_with_fmt_str(
92 move |c| formatted(c.as_materialized_series(), &format).map(Column::from),
93 text_field,
94 "past_calendar_format",
95 )
96}
97
98pub fn calendar_expr(expr: Expr) -> Expr {
101 expr.map_with_fmt_str(
102 |c| {
103 Ok(calendar_without_out_of_range(c.as_materialized_series())?
104 .map(Column::from)
105 .unwrap_or(c))
106 },
107 |_, field| Ok(field.clone()),
108 "past_calendar_null",
109 )
110}
111
112const NS_PER_DAY: i128 = 86_400_000_000_000;
113
114struct Reach {
117 back: i128,
118 forward: i128,
119 low: i128,
120}
121
122fn ns_reach(function: &TemporalFunction, zoned: bool, args: &[Column]) -> Option<Reach> {
130 let spans = || -> Option<Vec<(i128, bool, bool)>> {
132 let arg = args.first()?.unique().ok()?;
134 let text = arg.as_materialized_series().str().ok()?;
135 text.iter()
136 .flatten()
137 .map(|text| {
138 let d = Duration::try_parse(text).ok()?;
139 let span = i128::from(d.months().abs()) * 31 * NS_PER_DAY
140 + i128::from(d.weeks().abs()) * 7 * NS_PER_DAY
141 + i128::from(d.days().abs()) * NS_PER_DAY
142 + i128::from(d.nanoseconds().abs());
143 Some((span, d.negative(), d.months() != 0))
144 })
145 .collect()
146 };
147 let longest = |spans: &[(i128, bool, bool)], back: Option<bool>| {
148 spans
149 .iter()
150 .filter(|(_, negative, _)| back.is_none_or(|back| *negative == back))
151 .map(|(span, _, _)| *span)
152 .max()
153 .unwrap_or(0)
154 };
155 let in_months = |spans: &[(i128, bool, bool)]| spans.iter().any(|(_, _, months)| *months);
156 let (back, forward, calendar) = match function {
157 TemporalFunction::MonthStart => (31 * NS_PER_DAY, 0, true),
158 TemporalFunction::MonthEnd => (31 * NS_PER_DAY, 32 * NS_PER_DAY, true),
160 TemporalFunction::Truncate => {
161 let spans = spans()?;
162 (longest(&spans, None), 0, in_months(&spans))
163 }
164 TemporalFunction::Round => {
166 let spans = spans()?;
167 let every = longest(&spans, None);
168 (every, every, in_months(&spans))
169 }
170 #[cfg(feature = "sql")]
172 TemporalFunction::OffsetBy => {
173 let spans = spans()?;
174 let back = longest(&spans, Some(true));
175 (back, longest(&spans, Some(false)), in_months(&spans))
176 }
177 TemporalFunction::Date
179 | TemporalFunction::Time
180 | TemporalFunction::OrdinalDay
181 | TemporalFunction::IsoYear
182 | TemporalFunction::IsLeapYear
183 | TemporalFunction::DaysInMonth
184 | TemporalFunction::Datetime
185 if zoned =>
186 {
187 (0, 0, true)
188 }
189 _ => return None,
190 };
191 let zone = if zoned { NS_PER_DAY } else { 0 };
192 Some(Reach {
193 back: back + zone,
194 forward: forward + zone,
195 low: if calendar || zoned {
196 -9_223_372_036_000_000_000
197 } else {
198 i128::from(i64::MIN)
199 },
200 })
201}
202
203fn ns_within_reach(function: &TemporalFunction, cols: &mut [Column]) -> PolarsResult<Column> {
208 let value = std::mem::take(&mut cols[0]);
209 let DataType::Datetime(TimeUnit::Nanoseconds, zone) = value.dtype() else {
210 return Ok(value);
211 };
212 let Some(reach) = ns_reach(function, zone.is_some(), &cols[1..]) else {
213 return Ok(value);
214 };
215 let fits = |v: i64| {
216 i128::from(v) - reach.back >= reach.low
217 && i128::from(v) + reach.forward <= i128::from(i64::MAX)
218 };
219 let series = value.as_materialized_series();
220 let stored = series.to_physical_repr();
221 let stored = stored.i64()?;
222 if [stored.min(), stored.max()].into_iter().flatten().all(fits) {
224 return Ok(value);
225 }
226 let kept = stored.apply(|v| v.filter(|v| fits(*v)));
227 Ok(kept
228 .into_series()
229 .cast(value.dtype())?
230 .with_name(series.name().clone())
231 .into_column())
232}
233
234fn ns_edge_expr(input: &[Expr], function: TemporalFunction) -> Expr {
238 input[0].clone().map_many(
239 move |cols| ns_within_reach(&function, cols),
240 &input[1..],
241 |_, fields| Ok(fields[0].clone()),
242 )
243}
244
245fn moves_ns(function: &TemporalFunction) -> bool {
248 #[cfg(feature = "sql")]
249 if matches!(function, TemporalFunction::OffsetBy) {
250 return true;
251 }
252 matches!(
253 function,
254 TemporalFunction::MonthStart
255 | TemporalFunction::MonthEnd
256 | TemporalFunction::Truncate
257 | TemporalFunction::Round
258 | TemporalFunction::Date
259 | TemporalFunction::Time
260 | TemporalFunction::OrdinalDay
261 | TemporalFunction::IsoYear
262 | TemporalFunction::IsLeapYear
263 | TemporalFunction::DaysInMonth
264 | TemporalFunction::Datetime
265 )
266}
267
268fn countable_dates(column: Column, unit: TimeUnit) -> PolarsResult<Column> {
274 if column.dtype() != &DataType::Date {
275 return Ok(column);
276 }
277 let per_day: i64 = match unit {
278 TimeUnit::Nanoseconds => 86_400_000_000_000,
279 TimeUnit::Microseconds => 86_400_000_000,
280 TimeUnit::Milliseconds => 86_400_000,
281 };
282 let fits = |days: i32| i64::from(days).abs() <= i64::MAX / per_day;
283 let series = column.as_materialized_series();
284 let days = series.to_physical_repr();
285 let days = days.i32()?;
286 if [days.min(), days.max()].into_iter().flatten().all(fits) {
287 return Ok(column);
288 }
289 Ok(days
290 .apply(|days| days.filter(|days| fits(*days)))
291 .into_date()
292 .into_series()
293 .with_name(series.name().clone())
294 .into_column())
295}
296
297fn countable_expr(expr: Expr, unit: TimeUnit) -> Expr {
300 expr.map_with_fmt_str(
301 move |c| countable_dates(c, unit),
302 |_, field| Ok(field.clone()),
303 "past_calendar_countable",
304 )
305}
306
307fn reads_calendar(function: &TemporalFunction) -> bool {
310 !matches!(
311 function,
312 TemporalFunction::TimeStamp(_)
313 | TemporalFunction::CastTimeUnit(_)
314 | TemporalFunction::WithTimeUnit(_)
315 | TemporalFunction::ConvertTimeZone(_)
316 )
317}
318
319pub fn guard_expr(expr: Expr, schema: Option<&Schema>) -> Expr {
333 let may_leave = |e: &Expr| match (e, schema) {
334 (Expr::Literal(_), _) => false,
335 (e, Some(schema)) => e
336 .to_field(schema)
337 .map_or(true, |f| can_leave_calendar(f.dtype())),
338 (_, None) => true,
339 };
340 let may_be_ns = |e: &Expr| match (e, schema) {
341 (Expr::Literal(_), _) => false,
342 (e, Some(schema)) => e.to_field(schema).map_or(true, |f| {
343 matches!(f.dtype(), DataType::Datetime(TimeUnit::Nanoseconds, _))
344 }),
345 (_, None) => true,
346 };
347 let as_text = |e: Expr| {
348 if may_leave(&e) {
349 text_expr(e, CastOptions::NonStrict)
350 } else {
351 e
352 }
353 };
354 let gives_text = |e: &Expr| {
355 schema.is_some_and(|schema| {
356 e.to_field(schema)
357 .is_ok_and(|f| f.dtype() == &DataType::String)
358 })
359 };
360 let may_be_date = |e: &Expr| match (e, schema) {
361 (Expr::Literal(_), _) => false,
362 (e, Some(schema)) => e
363 .to_field(schema)
364 .map_or(true, |f| f.dtype() == &DataType::Date),
365 (_, None) => true,
366 };
367 let countable = |e: Expr, unit: TimeUnit| {
368 if may_be_date(&e) {
369 countable_expr(e, unit)
370 } else {
371 e
372 }
373 };
374 let datetime_unit = |e: &Expr| match e.to_field(schema?).ok()?.dtype() {
375 DataType::Datetime(unit, _) => Some(*unit),
376 _ => None,
377 };
378 expr.map_expr(|e| match e {
379 e @ (Expr::Ternary { .. }
380 | Expr::Function {
381 function: FunctionExpr::Coalesce,
382 ..
383 }) if gives_text(&e) => match e {
384 Expr::Ternary {
385 predicate,
386 truthy,
387 falsy,
388 } => Expr::Ternary {
389 predicate,
390 truthy: Arc::new(as_text(Arc::unwrap_or_clone(truthy))),
391 falsy: Arc::new(as_text(Arc::unwrap_or_clone(falsy))),
392 },
393 Expr::Function { input, function } => Expr::Function {
394 input: input.into_iter().map(as_text).collect(),
395 function,
396 },
397 e => e,
398 },
399 e @ (Expr::Ternary { .. }
401 | Expr::Function {
402 function:
403 FunctionExpr::Coalesce
404 | FunctionExpr::FillNull
405 | FunctionExpr::MaxHorizontal
406 | FunctionExpr::MinHorizontal,
407 ..
408 }) => match (datetime_unit(&e), e) {
409 (
410 Some(unit),
411 Expr::Ternary {
412 predicate,
413 truthy,
414 falsy,
415 },
416 ) => Expr::Ternary {
417 predicate,
418 truthy: Arc::new(countable(Arc::unwrap_or_clone(truthy), unit)),
419 falsy: Arc::new(countable(Arc::unwrap_or_clone(falsy), unit)),
420 },
421 (Some(unit), Expr::Function { input, function }) => Expr::Function {
422 input: input.into_iter().map(|e| countable(e, unit)).collect(),
423 function,
424 },
425 (_, e) => e,
426 },
427 Expr::Cast {
428 expr,
429 dtype,
430 options,
431 } => match dtype.as_literal() {
432 Some(DataType::String) if may_leave(&expr) => {
433 text_expr(Arc::unwrap_or_clone(expr), options)
434 }
435 Some(DataType::Datetime(unit, _)) if may_be_date(&expr) => Expr::Cast {
436 expr: Arc::new(countable_expr(Arc::unwrap_or_clone(expr), *unit)),
437 dtype,
438 options,
439 },
440 _ => Expr::Cast {
441 expr,
442 dtype,
443 options,
444 },
445 },
446 Expr::Function {
447 mut input,
448 function: FunctionExpr::TemporalExpr(function),
449 } if input
450 .first()
451 .is_some_and(|e| may_leave(e) || (moves_ns(&function) && may_be_ns(e))) =>
452 {
453 if may_leave(&input[0]) {
454 if let TemporalFunction::ToString(format) = function {
455 return format_expr(input.swap_remove(0), format);
456 }
457 if reads_calendar(&function) {
458 input[0] = calendar_expr(input[0].clone());
459 }
460 }
461 if moves_ns(&function) && may_be_ns(&input[0]) {
462 input[0] = ns_edge_expr(&input, function.clone());
463 }
464 Expr::Function {
465 input,
466 function: FunctionExpr::TemporalExpr(function),
467 }
468 }
469 #[cfg(feature = "sql")]
471 Expr::Function {
472 input,
473 function:
474 function @ FunctionExpr::StringExpr(
475 StringFunction::ConcatHorizontal { .. } | StringFunction::ConcatVertical { .. },
476 ),
477 } => Expr::Function {
478 input: input.into_iter().map(as_text).collect(),
479 function,
480 },
481 e => e,
482 })
483}
484
485#[cfg(feature = "sql")]
487fn guards(expr: &Expr) -> bool {
488 expr.into_iter().any(|e| match e {
489 Expr::Cast { dtype, .. } => matches!(
490 dtype.as_literal(),
491 Some(DataType::String | DataType::Datetime(..))
492 ),
493 Expr::Ternary { .. } => true,
494 Expr::Function { function, .. } => matches!(
495 function,
496 FunctionExpr::TemporalExpr(_)
497 | FunctionExpr::Coalesce
498 | FunctionExpr::FillNull
499 | FunctionExpr::MaxHorizontal
500 | FunctionExpr::MinHorizontal
501 | FunctionExpr::StringExpr(
502 StringFunction::ConcatHorizontal { .. } | StringFunction::ConcatVertical { .. }
503 )
504 ),
505 _ => false,
506 })
507}
508
509#[cfg(feature = "sql")]
511fn holds_guards(plan: &DslPlan) -> bool {
512 plan.into_iter().any(|node| match node {
513 DslPlan::Filter { predicate, .. } => guards(predicate),
514 DslPlan::Select { expr, .. } => expr.iter().any(guards),
515 DslPlan::HStack { exprs, .. } => exprs.iter().any(guards),
516 DslPlan::Sort { by_column, .. } => by_column.iter().any(guards),
517 DslPlan::GroupBy {
518 keys,
519 predicates,
520 aggs,
521 ..
522 } => keys.iter().chain(predicates).chain(aggs).any(guards),
523 DslPlan::Join {
524 left_on,
525 right_on,
526 predicates,
527 ..
528 } => left_on.iter().chain(right_on).chain(predicates).any(guards),
529 DslPlan::Union { args, .. } => args.to_supertypes,
531 _ => false,
532 })
533}
534
535#[cfg(feature = "sql")]
539fn union_text(input: &mut DslPlan, union: &Schema, diagonal: bool) {
540 let Ok(own) = LazyFrame::from(input.clone()).collect_schema() else {
541 return;
542 };
543 let texts: Vec<Expr> = own
544 .iter()
545 .enumerate()
546 .filter(|(i, (name, dtype))| {
547 let stacked = if diagonal {
548 union.get(name.as_str())
549 } else {
550 union.get_at_index(*i).map(|(_, dtype)| dtype)
551 };
552 can_leave_calendar(dtype) && stacked == Some(&DataType::String)
553 })
554 .map(|(_, (name, _))| text_expr(Expr::Column(name.clone()), CastOptions::NonStrict))
555 .collect();
556 if !texts.is_empty() {
557 *input = LazyFrame::from(input.clone())
558 .with_columns(texts)
559 .logical_plan;
560 }
561}
562
563#[cfg(feature = "sql")]
567pub fn guard_plan(plan: &mut DslPlan) {
568 if !holds_guards(plan) {
569 return;
570 }
571 let schema = |input: &DslPlan| LazyFrame::from(input.clone()).collect_schema().ok();
574 let guard = |exprs: &mut Vec<Expr>, schema: Option<&Schema>| {
575 for expr in exprs.iter_mut().filter(|e| guards(e)) {
576 *expr = guard_expr(std::mem::take(expr), schema);
577 }
578 };
579 match plan {
580 DslPlan::Filter { input, predicate } if guards(predicate) => {
581 *predicate = guard_expr(std::mem::take(predicate), schema(input).as_deref());
582 }
583 DslPlan::Select { input, expr, .. } if expr.iter().any(guards) => {
584 guard(expr, schema(input).as_deref());
585 }
586 DslPlan::HStack { input, exprs, .. } if exprs.iter().any(guards) => {
587 guard(exprs, schema(input).as_deref());
588 }
589 DslPlan::Sort {
590 input, by_column, ..
591 } if by_column.iter().any(guards) => {
592 guard(by_column, schema(input).as_deref());
593 }
594 DslPlan::GroupBy {
595 input,
596 keys,
597 predicates,
598 aggs,
599 ..
600 } if keys.iter().chain(&*predicates).chain(&*aggs).any(guards) => {
601 let schema = schema(input);
602 guard(keys, schema.as_deref());
603 guard(predicates, schema.as_deref());
604 guard(aggs, schema.as_deref());
605 }
606 DslPlan::Join {
607 input_left,
608 input_right,
609 left_on,
610 right_on,
611 predicates,
612 ..
613 } => {
614 if left_on.iter().any(guards) {
615 guard(left_on, schema(input_left).as_deref());
616 }
617 if right_on.iter().any(guards) {
618 guard(right_on, schema(input_right).as_deref());
619 }
620 guard(predicates, None);
622 }
623 DslPlan::Union { args, .. } if args.to_supertypes => {
624 let union = schema(plan);
625 if let (Some(union), DslPlan::Union { inputs, args }) = (union, &mut *plan) {
626 for input in inputs {
627 union_text(input, &union, args.diagonal);
628 }
629 }
630 }
631 _ => {}
632 }
633 if let DslPlan::IR { dsl, .. } = plan {
634 let mut inner = Arc::unwrap_or_clone(dsl.clone());
637 guard_plan(&mut inner);
638 *plan = inner;
639 return;
640 }
641 crate::widgets::datatable::for_each_input(plan, &mut guard_plan);
642}
643
644#[cfg(test)]
645mod tests {
646 use super::*;
647 use crate::nested_json::tests::calendar;
648
649 const IN_RANGE: usize = 7;
651
652 fn past_rows(series: &Series) -> Vec<Option<String>> {
654 series
655 .slice(IN_RANGE as i64, series.len() - IN_RANGE)
656 .iter()
657 .map(|v| (!v.is_null()).then(|| crate::exact::str_value(&v).into_owned()))
658 .collect()
659 }
660
661 #[test]
665 fn text_is_polars_own_and_a_date_past_the_calendar_its_stored_number() {
666 let in_range = calendar(false);
667 let past = calendar(true);
668 for column in past.columns() {
669 let series = column.as_materialized_series();
670 let polars = in_range
671 .column(column.name())
672 .unwrap()
673 .as_materialized_series();
674 for (text, expected) in [
675 (
676 cast_text(series, CastOptions::Strict).unwrap(),
677 polars.cast(&DataType::String).unwrap(),
678 ),
679 (
680 formatted(series, "%Y/%m/%d").unwrap(),
681 TemporalMethods::to_string(polars, "%Y/%m/%d").unwrap(),
682 ),
683 ] {
684 let name = column.name();
685 assert_eq!(text.dtype(), &DataType::String, "{name}");
686 assert_eq!(text.name(), name);
687 assert!(
688 text.head(Some(IN_RANGE)).equals_missing(&expected),
689 "{name}: {text:?}"
690 );
691 if can_leave_calendar(series.dtype()) {
692 let shown = past_rows(series);
693 assert!(
694 shown.iter().flatten().all(|s| s.contains("since")),
695 "{name}"
696 );
697 assert_eq!(past_rows(&text), shown, "{name}");
698 }
699 }
700 }
701 let numbers = Series::new("n".into(), [Some(1.5f64), None]);
703 assert!(
704 cast_text(&numbers, CastOptions::NonStrict)
705 .unwrap()
706 .equals_missing(&numbers.cast(&DataType::String).unwrap())
707 );
708 }
709
710 #[test]
713 fn date_parts_of_a_date_past_the_calendar_are_null() {
714 let in_range = calendar(false).lazy();
715 let past = calendar(true).lazy();
716 let schema = past.clone().collect_schema().unwrap();
717 for (name, dtype) in schema.iter() {
718 let c = || col(name.clone());
719 let mut parts = vec![
720 c().dt().date(),
721 c().dt().ordinal_day(),
722 c().dt().weekday(),
723 c().dt().iso_year(),
724 ];
725 if !matches!(dtype, DataType::Date) {
726 parts.push(c().dt().time());
727 }
728 parts.extend([
730 c().dt().month_start(),
731 c().dt().month_end(),
732 c().dt().truncate(lit("1mo")),
733 c().dt().round(lit("1mo")),
734 ]);
735 if matches!(dtype, DataType::Datetime(..)) {
736 parts.push(c().dt().truncate(lit("1d")));
737 }
738 for part in parts {
739 let guarded = guard_expr(part.clone(), Some(&schema));
740 let out = past.clone().select([guarded]).collect().unwrap();
741 let out = out.columns()[0].as_materialized_series();
742 let polars = in_range.clone().select([part.clone()]).collect().unwrap();
743 let case = format!("{name}: {part:?}");
744 assert!(
745 out.head(Some(IN_RANGE))
746 .equals_missing(polars.columns()[0].as_materialized_series()),
747 "{case}"
748 );
749 if can_leave_calendar(dtype) {
750 assert_eq!(past_rows(out), [None, None], "{case}");
751 }
752 }
753 }
754 }
755
756 #[test]
761 fn date_math_near_the_ends_of_the_nanosecond_range_is_null() {
762 const DAY: i64 = 86_400_000_000_000;
763 let stamps = [
766 Some(0),
767 Some(400 * DAY),
768 None,
769 Some(i64::MAX),
770 Some(i64::MIN + 1),
771 Some(i64::MAX - 20 * DAY),
772 Some(i64::MIN + DAY),
773 ];
774 let paris = TimeZone::opt_try_new(Some("Europe/Paris")).unwrap();
775 let frame = |rows: &[Option<i64>]| {
776 let at = |name: &str, zone: Option<TimeZone>| {
777 Series::new(name.into(), rows)
778 .cast(&DataType::Datetime(TimeUnit::Nanoseconds, zone))
779 .unwrap()
780 .into_column()
781 };
782 DataFrame::new_infer_height(vec![at("n", None), at("z", paris.clone())])
783 .unwrap()
784 .lazy()
785 };
786 let schema = frame(&stamps).collect_schema().unwrap();
787 let n = || col("n");
788 let z = || col("z");
789 let cases = [
791 (n().dt().month_start(), [true, false, true, false]),
792 (n().dt().month_end(), [false, false, false, false]),
794 (n().dt().truncate(lit("1d")), [true, false, true, true]),
795 (n().dt().truncate(lit("1mo")), [true, false, true, false]),
796 (n().dt().round(lit("1h")), [false, false, true, true]),
797 #[cfg(feature = "sql")]
798 (n().dt().offset_by(lit("1d")), [false, true, true, true]),
799 #[cfg(feature = "sql")]
800 (n().dt().offset_by(lit("-1mo")), [true, false, true, false]),
801 (n().dt().date(), [true, true, true, true]),
802 (n().dt().ordinal_day(), [true, true, true, true]),
803 (z().dt().month_start(), [false, false, true, false]),
804 (z().dt().month_end(), [false, false, false, false]),
805 #[cfg(feature = "sql")]
806 (z().dt().offset_by(lit("1d")), [false, false, true, false]),
807 (z().dt().date(), [false, false, true, false]),
808 (z().dt().time(), [false, false, true, false]),
809 (z().dt().ordinal_day(), [false, false, true, false]),
810 (z().dt().iso_year(), [false, false, true, false]),
811 (z().dt().datetime(), [false, false, true, false]),
812 (z().dt().year(), [true, true, true, true]),
813 ];
814 for (part, kept) in cases {
815 let guarded = guard_expr(part.clone(), Some(&schema));
816 let polars_alone = |v: Option<i64>| {
817 frame(&[v])
818 .select([part.clone()])
819 .collect()
820 .unwrap()
821 .columns()[0]
822 .get(0)
823 .unwrap()
824 .into_static()
825 };
826 let expected: Vec<AnyValue> = stamps
827 .iter()
828 .enumerate()
829 .map(|(i, v)| match i.checked_sub(3) {
830 Some(edge) if !kept[edge] => AnyValue::Null,
831 _ => polars_alone(*v),
832 })
833 .collect();
834 for streaming in [false, true] {
835 let out = crate::statistics::collect_lazy(
836 frame(&stamps).select([guarded.clone()]),
837 streaming,
838 )
839 .unwrap();
840 let out: Vec<AnyValue> = out.columns()[0]
841 .as_materialized_series()
842 .iter()
843 .map(|v| v.into_static())
844 .collect();
845 assert_eq!(out, expected, "{part:?} streaming: {streaming}");
846 }
847 }
848 }
849
850 #[test]
853 fn only_operations_on_dates_change() {
854 let schema = Schema::from_iter([
855 Field::new("i".into(), DataType::Int64),
856 Field::new("s".into(), DataType::String),
857 Field::new("ns".into(), DataType::Datetime(TimeUnit::Nanoseconds, None)),
858 Field::new("d".into(), DataType::Date),
859 Field::new("t".into(), DataType::Datetime(TimeUnit::Microseconds, None)),
860 ]);
861 let kept = [
862 col("i").cast(DataType::String),
863 col("s").cast(DataType::String).str().len_chars(),
864 col("ns").cast(DataType::String),
865 col("ns").dt().year(),
866 lit(1).cast(DataType::String),
867 col("t").dt().timestamp(TimeUnit::Milliseconds),
868 col("d").cast(DataType::Int64),
869 #[cfg(feature = "sql")]
870 concat_str([col("i"), col("s")], "", true),
871 ];
872 for expr in kept {
873 assert_eq!(guard_expr(expr.clone(), Some(&schema)), expr);
874 }
875 let changed = [
876 col("d").cast(DataType::String),
877 col("t").cast(DataType::String).str().len_chars(),
878 col("t").max().cast(DataType::String),
879 col("t").dt().to_string("%Y"),
880 col("d").dt().month_start(),
881 col("ns").dt().month_start(),
882 #[cfg(feature = "sql")]
883 concat_str([col("i"), col("t")], "", true),
884 ];
885 for expr in changed {
886 assert_ne!(guard_expr(expr.clone(), Some(&schema)), expr);
887 assert_ne!(guard_expr(expr.clone(), None), expr);
888 }
889 assert_ne!(
890 guard_expr(col("i").cast(DataType::String), None),
891 col("i").cast(DataType::String)
892 );
893 let one = || col("i").eq(lit(1));
896 let kept = [
897 coalesce(&[col("t"), col("ns")]),
898 coalesce(&[col("s"), lit("x")]),
899 when(one()).then(col("t")).otherwise(col("ns")),
900 when(one()).then(col("ns")).otherwise(col("s")),
901 col("t").cast(DataType::Datetime(TimeUnit::Nanoseconds, None)),
902 col("d").fill_null(col("d")),
903 ];
904 for expr in kept {
905 assert_eq!(guard_expr(expr.clone(), Some(&schema)), expr);
906 }
907 let changed = [
908 coalesce(&[col("t"), lit("x")]),
909 coalesce(&[col("s"), col("d")]),
910 when(one()).then(col("d")).otherwise(col("s")),
911 when(one()).then(lit("x")).otherwise(col("t")),
912 coalesce(&[col("d"), col("t")]),
914 when(one()).then(col("d")).otherwise(col("t")),
915 col("d").fill_null(col("t")),
916 polars::lazy::dsl::max_horizontal([col("t"), col("d")]).unwrap(),
917 ];
918 for expr in changed {
919 assert_ne!(guard_expr(expr.clone(), Some(&schema)), expr);
920 assert_eq!(guard_expr(expr.clone(), None), expr);
921 }
922 let cast = col("d").cast(DataType::Datetime(TimeUnit::Microseconds, None));
924 assert_ne!(guard_expr(cast.clone(), Some(&schema)), cast);
925 assert_ne!(guard_expr(cast.clone(), None), cast);
926 }
927
928 #[cfg(feature = "sql")]
932 #[test]
933 fn a_date_met_with_text_is_its_text() {
934 let frame = |past: bool| {
935 let at = if past { i64::MIN + 1 } else { 1 };
936 df!(
937 "s" => ["a", "b"],
938 "d" => [0, if past { i32::MAX } else { 1 }],
939 "t" => [0, at],
940 )
941 .unwrap()
942 .lazy()
943 .with_columns([
944 col("d").cast(DataType::Date),
945 col("t").cast(DataType::Datetime(TimeUnit::Milliseconds, None)),
946 ])
947 };
948 let mut ctx = polars_sql::SQLContext::new();
949 for sql in [
950 "SELECT COALESCE(t, 'x') AS x FROM df",
951 "SELECT CASE WHEN s = 'b' THEN d ELSE s END AS x FROM df",
952 "SELECT d AS x FROM df UNION ALL SELECT s FROM df",
953 "SELECT s AS x FROM df UNION SELECT t FROM df",
954 ] {
955 for past in [false, true] {
956 ctx.register("df", frame(past));
957 let polars = ctx.execute(sql).unwrap();
958 let mut guarded = polars.clone();
959 guard_plan(&mut guarded.logical_plan);
960 let values = |lf: LazyFrame| {
962 let df = lf.collect().unwrap();
963 let mut values: Vec<Option<String>> = df
964 .column("x")
965 .unwrap()
966 .str()
967 .unwrap()
968 .iter()
969 .map(|v| v.map(str::to_string))
970 .collect();
971 values.sort();
972 values
973 };
974 let text = values(guarded);
975 if past {
976 assert!(text.iter().flatten().any(|s| s.contains("since")), "{sql}");
977 } else {
978 assert_eq!(text, values(polars), "{sql}");
979 }
980 }
981 }
982 }
983
984 const DAYS: [Option<i32>; 13] = [
987 Some(0),
988 Some(19_737),
989 Some(106_751),
990 Some(106_752),
991 Some(-106_751),
992 Some(-106_752),
993 Some(106_751_991),
994 Some(106_751_992),
995 Some(-106_751_991),
996 Some(-106_751_992),
997 Some(i32::MAX),
998 Some(i32::MIN),
999 None,
1000 ];
1001
1002 fn dates_and_datetimes(days: &[Option<i32>]) -> LazyFrame {
1005 let n = days.len() as i64;
1006 df!(
1007 "d" => days,
1008 "t" => (0..n).map(|i| (i % 3 != 0).then_some(i * 1_000_000)).collect::<Vec<_>>(),
1009 "b" => (0..n).map(|i| i % 2 == 0).collect::<Vec<_>>(),
1010 )
1011 .unwrap()
1012 .lazy()
1013 .with_columns([
1014 col("d").cast(DataType::Date),
1015 col("t").cast(DataType::Datetime(TimeUnit::Microseconds, None)),
1016 col("t")
1017 .cast(DataType::Datetime(TimeUnit::Nanoseconds, None))
1018 .alias("tn"),
1019 ])
1020 }
1021
1022 fn countable_days(unit: TimeUnit) -> Vec<Option<i32>> {
1024 let per_day: i64 = match unit {
1025 TimeUnit::Nanoseconds => 86_400_000_000_000,
1026 _ => 86_400_000_000,
1027 };
1028 DAYS.iter()
1029 .map(|d| d.filter(|d| i64::from(*d).abs() <= i64::MAX / per_day))
1030 .collect()
1031 }
1032
1033 #[test]
1038 fn a_date_a_datetime_cannot_count_is_null_as_one() {
1039 use polars::lazy::dsl::{max_horizontal, min_horizontal};
1040 let schema = dates_and_datetimes(&DAYS).collect_schema().unwrap();
1041 let us = TimeUnit::Microseconds;
1042 let ns = TimeUnit::Nanoseconds;
1043 let paris = TimeZone::opt_try_new(Some("Europe/Paris")).unwrap();
1044 let cases = [
1045 (col("d").strict_cast(DataType::Datetime(us, None)), us),
1046 (col("d").strict_cast(DataType::Datetime(ns, None)), ns),
1047 (col("d").strict_cast(DataType::Datetime(us, paris)), us),
1048 (coalesce(&[col("d"), col("t")]), us),
1049 (coalesce(&[col("t"), col("d")]), us),
1050 (coalesce(&[col("d"), col("tn")]), ns),
1051 (when(col("b")).then(col("d")).otherwise(col("t")), us),
1052 (max_horizontal([col("d"), col("t")]).unwrap(), us),
1053 (min_horizontal([col("t"), col("d")]).unwrap(), us),
1054 (col("d").fill_null(col("t")), us),
1055 ];
1056 for (expr, unit) in cases {
1057 let guarded = guard_expr(expr.clone(), Some(&schema));
1058 let expected = dates_and_datetimes(&countable_days(unit))
1059 .select([expr.clone()])
1060 .collect()
1061 .unwrap();
1062 for streaming in [false, true] {
1063 let out = crate::statistics::collect_lazy(
1064 dates_and_datetimes(&DAYS).select([guarded.clone()]),
1065 streaming,
1066 )
1067 .unwrap();
1068 assert_eq!(out.schema(), expected.schema(), "{expr:?}");
1069 assert!(
1070 out.equals_missing(&expected),
1071 "{expr:?} streaming: {streaming}\n{out}\n{expected}"
1072 );
1073 }
1074 }
1075 }
1076
1077 #[cfg(feature = "sql")]
1080 #[test]
1081 fn a_date_a_timestamp_cannot_count_is_null_in_sql() {
1082 let mut ctx = polars_sql::SQLContext::new();
1083 for (sql, unit) in [
1084 (
1085 "SELECT CAST(d AS TIMESTAMP) AS x FROM df",
1086 TimeUnit::Microseconds,
1087 ),
1088 ("SELECT d::timestamp AS x FROM df", TimeUnit::Microseconds),
1089 (
1090 "SELECT CAST(d AS TIMESTAMP(9)) AS x FROM df",
1091 TimeUnit::Nanoseconds,
1092 ),
1093 ("SELECT COALESCE(d, t) AS x FROM df", TimeUnit::Microseconds),
1094 ("SELECT IFNULL(t, d) AS x FROM df", TimeUnit::Microseconds),
1095 ("SELECT COALESCE(d, tn) AS x FROM df", TimeUnit::Nanoseconds),
1096 (
1097 "SELECT CASE WHEN b THEN d ELSE t END AS x FROM df",
1098 TimeUnit::Microseconds,
1099 ),
1100 ("SELECT GREATEST(d, t) AS x FROM df", TimeUnit::Microseconds),
1101 ("SELECT LEAST(t, d) AS x FROM df", TimeUnit::Microseconds),
1102 (
1103 "SELECT MAX(CAST(d AS TIMESTAMP)) AS x FROM df GROUP BY b",
1104 TimeUnit::Microseconds,
1105 ),
1106 (
1107 "SELECT df.d AS x FROM df JOIN e ON CAST(df.d AS TIMESTAMP) = e.t",
1108 TimeUnit::Microseconds,
1109 ),
1110 ] {
1111 let mut run = |days: &[Option<i32>], guard: bool| {
1113 ctx.register("df", dates_and_datetimes(days));
1114 ctx.register("e", dates_and_datetimes(days));
1115 let mut lf = ctx.execute(sql).unwrap();
1116 if guard {
1117 guard_plan(&mut lf.logical_plan);
1118 }
1119 lf.sort(["x"], SortMultipleOptions::default())
1120 .collect()
1121 .unwrap()
1122 };
1123 let expected = run(&countable_days(unit), false);
1124 let out = run(&DAYS, true);
1125 assert!(out.equals_missing(&expected), "{sql}\n{out}\n{expected}");
1126 }
1127 }
1128}