diff --git a/crates/polars-expr/src/dispatch/strings.rs b/crates/polars-expr/src/dispatch/strings.rs index d8792f6334c5..99de86c8dbcf 100644 --- a/crates/polars-expr/src/dispatch/strings.rs +++ b/crates/polars-expr/src/dispatch/strings.rs @@ -466,23 +466,39 @@ pub(super) fn split_regex(s: &[Column], inclusive: bool, strict: bool) -> Polars Ok(out.into_column()) } +/// Yield nulls when a non-strict parse could not infer a format, and report the +/// values that failed when strict. +#[cfg(feature = "temporal")] +fn finish_strptime( + out: PolarsResult, + input: &Column, + dtype: &DataType, + options: &StrptimeOptions, +) -> PolarsResult { + let out = match out { + Err(_) if !options.strict && options.format.is_none() => { + Column::full_null(input.name().clone(), input.len(), dtype) + }, + out => out?, + }; + + if options.strict && input.null_count() != out.null_count() { + handle_casting_failures(input.as_materialized_series(), out.as_materialized_series())?; + } + Ok(out) +} + #[cfg(feature = "dtype-date")] fn to_date(s: &Column, options: &StrptimeOptions) -> PolarsResult { let ca = s.str()?; - let out = { - if options.exact { - ca.as_date(options.format.as_deref(), options.cache)? - .into_column() - } else { - ca.as_date_not_exact(options.format.as_deref())? - .into_column() - } - }; - - if options.strict && ca.null_count() != out.null_count() { - handle_casting_failures(s.as_materialized_series(), out.as_materialized_series())?; + let out = if options.exact { + ca.as_date(options.format.as_deref(), options.cache) + } else { + ca.as_date_not_exact(options.format.as_deref()) } - Ok(out.into_column()) + .map(|ca| ca.into_column()); + + finish_strptime(out, s, &DataType::Date, options) } #[cfg(feature = "dtype-datetime")] @@ -511,33 +527,28 @@ fn to_datetime( }; let out = if options.exact { - datetime_strings - .as_datetime( - options.format.as_deref(), - *time_unit, - options.cache, - tz_aware, - time_zone, - ambiguous, - )? - .into_column() + datetime_strings.as_datetime( + options.format.as_deref(), + *time_unit, + options.cache, + tz_aware, + time_zone, + ambiguous, + ) } else { - datetime_strings - .as_datetime_not_exact( - options.format.as_deref(), - *time_unit, - tz_aware, - time_zone, - ambiguous, - true, - )? - .into_column() - }; - - if options.strict && datetime_strings.null_count() != out.null_count() { - handle_casting_failures(s[0].as_materialized_series(), out.as_materialized_series())?; + datetime_strings.as_datetime_not_exact( + options.format.as_deref(), + *time_unit, + tz_aware, + time_zone, + ambiguous, + true, + ) } - Ok(out.into_column()) + .map(|ca| ca.into_column()); + + let dtype = DataType::Datetime(*time_unit, time_zone.cloned()); + finish_strptime(out, &s[0], &dtype, options) } #[cfg(feature = "dtype-time")] @@ -546,15 +557,12 @@ fn to_time(s: &Column, options: &StrptimeOptions) -> PolarsResult { options.exact, ComputeError: "non-exact not implemented for Time data type" ); - let ca = s.str()?; - let out = ca - .as_time(options.format.as_deref(), options.cache)? - .into_column(); + let out = s + .str()? + .as_time(options.format.as_deref(), options.cache) + .map(|ca| ca.into_column()); - if options.strict && ca.null_count() != out.null_count() { - handle_casting_failures(s.as_materialized_series(), out.as_materialized_series())?; - } - Ok(out.into_column()) + finish_strptime(out, s, &DataType::Time, options) } #[cfg(feature = "concat_str")] diff --git a/crates/polars-sql/src/sql_expr.rs b/crates/polars-sql/src/sql_expr.rs index d3a2850c3787..d9ab239b12c7 100644 --- a/crates/polars-sql/src/sql_expr.rs +++ b/crates/polars-sql/src/sql_expr.rs @@ -116,6 +116,23 @@ fn cast_literal_series(s: &Series, dtype: &DataType) -> PolarsResult { s.strict_cast(dtype) } +/// Parse a string expression into `Date`/`Time`/`Datetime`; `None` for other dtypes. +fn parse_string_as_temporal(expr: Expr, dtype: &DataType, strict: bool) -> Option { + let options = StrptimeOptions { + strict, + ..Default::default() + }; + Some(match dtype { + DataType::Date => expr.str().to_date(options), + DataType::Time => expr.str().to_time(options), + DataType::Datetime(tu, tz) => { + expr.str() + .to_datetime(Some(*tu), tz.clone(), options, lit("latest")) + }, + _ => return None, + }) +} + /// Extract the literal value; returns `(sql_value, optional_op)`. fn extract_literal_with_op<'a>( expr: &'a SQLExpr, @@ -550,39 +567,27 @@ impl SQLExprVisitor<'_> { }, |dt| dt.as_literal(), ); - match left_dtype { - Some(DataType::Time) if is_iso_time(s) => { - right.clone().str().to_time(StrptimeOptions { - strict: true, - ..Default::default() - }) + let parsed = match left_dtype { + Some(dtype @ DataType::Time) if is_iso_time(s) => { + parse_string_as_temporal(right.clone(), dtype, true) }, - Some(DataType::Date) if is_iso_date(s) => { - right.clone().str().to_date(StrptimeOptions { - strict: true, - ..Default::default() - }) + Some(dtype @ DataType::Date) if is_iso_date(s) => { + parse_string_as_temporal(right.clone(), dtype, true) }, - Some(DataType::Datetime(tu, tz)) if is_iso_datetime(s) || is_iso_date(s) => { - if s.len() == 10 { + Some(dtype @ DataType::Datetime(_, _)) + if is_iso_datetime(s) || is_iso_date(s) => + { + let s = if s.len() == 10 { // handle upcast from ISO date string (10 chars) to datetime - lit(format!("{s}T00:00:00")) + format!("{s}T00:00:00") } else { - lit(s.replacen(' ', "T", 1)) - } - .str() - .to_datetime( - Some(*tu), - tz.clone(), - StrptimeOptions { - strict: true, - ..Default::default() - }, - lit("latest"), - ) + s.replacen(' ', "T", 1) + }; + parse_string_as_temporal(lit(s), dtype, true) }, - _ => right.clone(), - } + _ => None, + }; + parsed.unwrap_or_else(|| right.clone()) } } else { right.clone() @@ -954,7 +959,7 @@ impl SQLExprVisitor<'_> { dtype, DataType::Date | DataType::Time | DataType::Datetime(_, _) ) { - return elems.strict_cast(dtype); + return cast_literal_series(&elems, dtype); } } } @@ -1005,12 +1010,31 @@ impl SQLExprVisitor<'_> { return Ok(expr.str().json_decode(DataType::Struct(Vec::new()))); } let polars_type = map_sql_dtype_to_polars(dtype)?; - Ok(match cast_kind { - CastKind::Cast | CastKind::DoubleColon => expr.strict_cast(polars_type), - CastKind::TryCast | CastKind::SafeCast => expr.cast(polars_type), + let strict = matches!(cast_kind, CastKind::Cast | CastKind::DoubleColon); + + // `CAST( AS DATE/TIME/TIMESTAMP)` parses rather than casts + if matches!( + polars_type, + DataType::Date | DataType::Time | DataType::Datetime(_, _) + ) && self.is_string_expr(&expr) + && let Some(parsed) = parse_string_as_temporal(expr.clone(), &polars_type, strict) + { + return Ok(parsed); + } + Ok(if strict { + expr.strict_cast(polars_type) + } else { + expr.cast(polars_type) }) } + /// Whether `expr` is known to be `String`; false if the dtype cannot be resolved. + fn is_string_expr(&self, expr: &Expr) -> bool { + let empty = Schema::default(); + let schema = self.active_schema.unwrap_or(&empty); + matches!(expr.to_field(schema), Ok(fld) if fld.dtype == DataType::String) + } + /// Visit a SQL literal. /// /// e.g. 1, 'foo', 1.0, NULL diff --git a/crates/polars-stream/src/nodes/strptime_infer.rs b/crates/polars-stream/src/nodes/strptime_infer.rs index 01c243aa3e8c..31aad13c1537 100644 --- a/crates/polars-stream/src/nodes/strptime_infer.rs +++ b/crates/polars-stream/src/nodes/strptime_infer.rs @@ -4,7 +4,7 @@ use polars_time::chunkedarray::StringMethods; use polars_time::chunkedarray::string::Pattern; use polars_time::chunkedarray::string::infer::{ DatetimeInfer, TryFromWithUnit, coerce_string_to_date, coerce_string_to_datetime, - infer_pattern_date_single, infer_pattern_datetime_single, sniff_time_fmt, + infer_from_values, infer_pattern_date_single, infer_pattern_datetime_single, sniff_time_fmt, }; use super::compute_node_prelude::*; @@ -15,6 +15,9 @@ pub struct StrptimeInferNode { infer: Option, phase: Phase, + /// Name to report parse failures against; the input carries an internal one. + input_name: PlSmallStr, + /// Ambiguous can be `raise`, `earliest`, `latest` and `null`. /// /// If it broadcast and it is `raise` or `null`, we can actually execute it here. So @@ -43,12 +46,18 @@ enum FormatInfer { } impl StrptimeInferNode { - pub fn new(dtype: DataType, options: StrptimeOptions, ambiguous_is_raise: bool) -> Self { + pub fn new( + dtype: DataType, + options: StrptimeOptions, + input_name: PlSmallStr, + ambiguous_is_raise: bool, + ) -> Self { Self { dtype, options, infer: None, phase: Phase::Inferring, + input_name, ambiguous_is_raise, } } @@ -56,14 +65,14 @@ impl StrptimeInferNode { impl FormatInfer { fn try_new( - val: &str, + ca: &StringChunked, dtype: &DataType, options: &StrptimeOptions, ) -> PolarsResult> { match dtype { #[cfg(feature = "dtype-date")] DataType::Date => { - let Some(pattern) = infer_pattern_date_single(val) else { + let Some(pattern) = infer_from_values(ca, infer_pattern_date_single) else { return Ok(None); }; let infer = DatetimeInfer::::try_from_with_unit(pattern, None)?; @@ -71,7 +80,7 @@ impl FormatInfer { }, #[cfg(feature = "dtype-datetime")] DataType::Datetime(tu, tz) => { - let Some(pattern) = infer_pattern_datetime_single(val) else { + let Some(pattern) = infer_from_values(ca, infer_pattern_datetime_single) else { return Ok(None); }; if matches!(pattern, Pattern::DatetimeYMDZ) && tz.is_none() { @@ -81,7 +90,10 @@ impl FormatInfer { Ok(Some(FormatInfer::Datetime(infer, tz.clone()))) }, #[cfg(feature = "dtype-time")] - DataType::Time => Ok(sniff_time_fmt(val).map(|f| FormatInfer::Time(f, options.cache))), + DataType::Time => { + Ok(infer_from_values(ca, sniff_time_fmt) + .map(|f| FormatInfer::Time(f, options.cache))) + }, _ => Ok(None), } } @@ -91,6 +103,7 @@ impl FormatInfer { col: &Column, ambiguous: &StringChunked, strict: bool, + input_name: &PlSmallStr, ) -> PolarsResult { let ca = col.str()?; let name = col.name().clone(); @@ -119,9 +132,11 @@ impl FormatInfer { }; if strict && col.null_count() != result.null_count() { + // reported against the name the user wrote, not the internal alias + let reported = result.clone().with_name(input_name.clone()); polars_core::utils::handle_casting_failures( col.as_materialized_series(), - result.as_materialized_series(), + reported.as_materialized_series(), )?; } @@ -174,15 +189,25 @@ impl ComputeNode for StrptimeInferNode { let dtype = &self.dtype; let options = &self.options; + let input_name = &self.input_name; let infer_slot = &mut self.infer; join_handles.push(scope.spawn_task(TaskPriority::High, async move { while let Ok(morsel) = recv.recv().await { if infer_slot.is_none() { let df = morsel.df().await; let ca = df.columns()[0].str()?; - if let Some(idx) = ca.first_non_null() { - *infer_slot = - FormatInfer::try_new(ca.get(idx).unwrap(), dtype, options)?; + if ca.null_count() != ca.len() { + *infer_slot = FormatInfer::try_new(ca, dtype, options)?; + + let unit = if matches!(dtype, DataType::Time) { + "time" + } else { + "date" + }; + polars_ensure!( + infer_slot.is_some() || !options.strict, + parse_fmt_idk = unit + ); } } @@ -196,7 +221,7 @@ impl ComputeNode for StrptimeInferNode { let cols = df.columns(); if let Some(ref mut infer) = *infer_slot { infer - .apply(&cols[0], &ambiguous, options.strict) + .apply(&cols[0], &ambiguous, options.strict, input_name) .map(Column::into_frame) } else { Ok(Column::full_null( @@ -221,6 +246,7 @@ impl ComputeNode for StrptimeInferNode { let senders = send_ports[0].take().unwrap().parallel(); for (mut recv, mut send) in receivers.into_iter().zip(senders) { let strict = self.options.strict; + let input_name = self.input_name.clone(); let ambiguous = ambiguous.clone(); let mut infer = self.infer.clone().unwrap(); join_handles.push(scope.spawn_task(TaskPriority::High, async move { @@ -229,7 +255,7 @@ impl ComputeNode for StrptimeInferNode { .try_map(|df| { let cols = df.columns(); infer - .apply(&cols[0], &ambiguous, strict) + .apply(&cols[0], &ambiguous, strict, &input_name) .map(Column::into_frame) }) .await?; diff --git a/crates/polars-stream/src/physical_plan/lower_expr.rs b/crates/polars-stream/src/physical_plan/lower_expr.rs index 9a9f298dd347..b5968b18bc2e 100644 --- a/crates/polars-stream/src/physical_plan/lower_expr.rs +++ b/crates/polars-stream/src/physical_plan/lower_expr.rs @@ -1565,6 +1565,7 @@ fn lower_exprs_with_ctx( } if options.format.is_none() && matches!(ctx.expr_arena.get(inner_exprs[1].node()), AExpr::Literal(s) if matches!(s.extract_str(), Some("raise" | "null"))) => { + let input_name = inner_exprs[0].output_name().clone(); let col_name = unique_column_name(); let select_stream = build_select_stream_with_ctx( input, @@ -1582,6 +1583,7 @@ fn lower_exprs_with_ctx( input: select_stream, dtype: dtype.as_ref().clone(), options: options.clone(), + input_name, ambiguous_is_raise, }; let node_key = ctx.phys_sm.insert(PhysNode::new(output_schema, kind)); diff --git a/crates/polars-stream/src/physical_plan/mod.rs b/crates/polars-stream/src/physical_plan/mod.rs index 26f28e933819..3ecf95c0c362 100644 --- a/crates/polars-stream/src/physical_plan/mod.rs +++ b/crates/polars-stream/src/physical_plan/mod.rs @@ -285,6 +285,9 @@ pub enum PhysNodeKind { dtype: DataType, options: StrptimeOptions, + /// Name the input had before lowering aliased it; used in error messages. + input_name: PlSmallStr, + /// Ambiguous can be `raise`, `earliest`, `latest` and `null`. /// /// If it is broadcast and it is `raise` or `null`, we can actually execute it in this diff --git a/crates/polars-stream/src/physical_plan/to_graph.rs b/crates/polars-stream/src/physical_plan/to_graph.rs index 5aaa6cb83c30..18e64ecce05c 100644 --- a/crates/polars-stream/src/physical_plan/to_graph.rs +++ b/crates/polars-stream/src/physical_plan/to_graph.rs @@ -561,6 +561,7 @@ fn to_graph_rec<'a>( input, dtype, options, + input_name, ambiguous_is_raise, } => { let input_key = to_graph_rec(input.node, ctx)?; @@ -568,6 +569,7 @@ fn to_graph_rec<'a>( nodes::strptime_infer::StrptimeInferNode::new( dtype.clone(), options.clone(), + input_name.clone(), *ambiguous_is_raise, ), [(input_key, input.port)], diff --git a/crates/polars-time/src/chunkedarray/string/infer.rs b/crates/polars-time/src/chunkedarray/string/infer.rs index d165033aa820..a4afe96d68eb 100644 --- a/crates/polars-time/src/chunkedarray/string/infer.rs +++ b/crates/polars-time/src/chunkedarray/string/infer.rs @@ -424,6 +424,11 @@ pub fn sniff_time_fmt(val: &str) -> Option<&'static str> { .find(|fmt| NaiveTime::parse_from_str(val, fmt).is_ok()) } +/// Scan the non-null values for the first that `infer` accepts. +pub fn infer_from_values(ca: &StringChunked, infer: impl FnMut(&str) -> Option) -> Option { + ca.iter().flatten().find_map(infer) +} + #[cfg(feature = "dtype-datetime")] pub fn to_datetime_with_inferred_tz( ca: &StringChunked, @@ -459,27 +464,23 @@ pub fn to_datetime( // Ensure that the inferred time_zone matches the given time_zone. ensure_matching_time_zone: bool, ) -> PolarsResult { - match ca.first_non_null() { - None => { - Ok(Int64Chunked::full_null(ca.name().clone(), ca.len()).into_datetime(tu, tz.cloned())) - }, - Some(idx) => { - let subset = ca.slice(idx as i64, ca.len()); - let pattern = subset - .iter() - .find_map(|opt_val| opt_val.and_then(infer_pattern_datetime_single)) - .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?; - let mut infer = DatetimeInfer::::try_from_with_unit(pattern, Some(tu))?; - #[cfg(feature = "timezones")] - if matches!(pattern, Pattern::DatetimeYMDZ) { - polars_ensure!( - !ensure_matching_time_zone || tz.is_some(), - to_datetime_tz_mismatch - ); - } - coerce_string_to_datetime(&mut infer, ca, tz, ambiguous) - }, + if ca.null_count() == ca.len() { + return Ok( + Int64Chunked::full_null(ca.name().clone(), ca.len()).into_datetime(tu, tz.cloned()) + ); } + + let pattern = infer_from_values(ca, infer_pattern_datetime_single) + .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?; + let mut infer = DatetimeInfer::::try_from_with_unit(pattern, Some(tu))?; + #[cfg(feature = "timezones")] + if matches!(pattern, Pattern::DatetimeYMDZ) { + polars_ensure!( + !ensure_matching_time_zone || tz.is_some(), + to_datetime_tz_mismatch + ); + } + coerce_string_to_datetime(&mut infer, ca, tz, ambiguous) } /// Apply a pre-built `DatetimeInfer` to a `StringChunked`, returning a `DateChunked`. #[cfg(feature = "dtype-date")] @@ -529,16 +530,12 @@ pub fn coerce_string_to_datetime( #[cfg(feature = "dtype-date")] pub(crate) fn to_date(ca: &StringChunked) -> PolarsResult { - match ca.first_non_null() { - None => Ok(Int32Chunked::full_null(ca.name().clone(), ca.len()).into_date()), - Some(idx) => { - let subset = ca.slice(idx as i64, ca.len()); - let pattern = subset - .iter() - .find_map(|opt_val| opt_val.and_then(infer_pattern_date_single)) - .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?; - let mut infer = DatetimeInfer::::try_from_with_unit(pattern, None).unwrap(); - coerce_string_to_date(&mut infer, ca) - }, + if ca.null_count() == ca.len() { + return Ok(Int32Chunked::full_null(ca.name().clone(), ca.len()).into_date()); } + + let pattern = infer_from_values(ca, infer_pattern_date_single) + .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))?; + let mut infer = DatetimeInfer::::try_from_with_unit(pattern, None).unwrap(); + coerce_string_to_date(&mut infer, ca) } diff --git a/crates/polars-time/src/chunkedarray/string/mod.rs b/crates/polars-time/src/chunkedarray/string/mod.rs index dbd683953e9d..2ece7ecc60c5 100644 --- a/crates/polars-time/src/chunkedarray/string/mod.rs +++ b/crates/polars-time/src/chunkedarray/string/mod.rs @@ -50,23 +50,6 @@ where .copied() } -#[cfg(feature = "dtype-datetime")] -fn sniff_fmt_datetime(val: &str) -> PolarsResult<&'static str> { - datetime_pattern(val, NaiveDateTime::parse_from_str) - .or_else(|| datetime_pattern(val, NaiveDate::parse_from_str)) - .ok_or_else(|| polars_err!(parse_fmt_idk = "datetime")) -} - -#[cfg(feature = "dtype-date")] -fn sniff_fmt_date(val: &str) -> PolarsResult<&'static str> { - date_pattern(val, NaiveDate::parse_from_str).ok_or_else(|| polars_err!(parse_fmt_idk = "date")) -} - -#[cfg(feature = "dtype-time")] -fn sniff_fmt_time(val: &str) -> PolarsResult<&'static str> { - time_pattern(val, NaiveTime::parse_from_str).ok_or_else(|| polars_err!(parse_fmt_idk = "time")) -} - pub trait StringMethods: AsString { #[cfg(feature = "dtype-time")] /// Parsing string values and return a [`TimeChunked`] @@ -75,14 +58,16 @@ pub trait StringMethods: AsString { let fmt = match fmt { Some(fmt) => fmt, None => { - let Some(idx) = string_ca.first_non_null() else { + if string_ca.null_count() == string_ca.len() { return Ok( Int64Chunked::full_null(string_ca.name().clone(), string_ca.len()) .into_time(), ); - }; - let val = string_ca.get(idx).expect("should not be null"); - sniff_fmt_time(val)? + } + infer::infer_from_values(string_ca, |val| { + time_pattern(val, NaiveTime::parse_from_str) + }) + .ok_or_else(|| polars_err!(parse_fmt_idk = "time"))? }, }; let use_cache = use_cache && string_ca.len() > 50; @@ -107,14 +92,16 @@ pub trait StringMethods: AsString { let fmt = match fmt { Some(fmt) => fmt, None => { - let Some(idx) = string_ca.first_non_null() else { + if string_ca.null_count() == string_ca.len() { return Ok( Int32Chunked::full_null(string_ca.name().clone(), string_ca.len()) .into_date(), ); - }; - let val = string_ca.get(idx).expect("should not be null"); - sniff_fmt_date(val)? + } + infer::infer_from_values(string_ca, |val| { + date_pattern(val, NaiveDate::parse_from_str) + }) + .ok_or_else(|| polars_err!(parse_fmt_idk = "date"))? }, }; let ca = unary_elementwise(string_ca, |opt_s| { @@ -154,14 +141,17 @@ pub trait StringMethods: AsString { let fmt = match fmt { Some(fmt) => fmt, None => { - let Some(idx) = string_ca.first_non_null() else { + if string_ca.null_count() == string_ca.len() { return Ok( Int64Chunked::full_null(string_ca.name().clone(), string_ca.len()) .into_datetime(tu, tz.cloned()), ); - }; - let val = string_ca.get(idx).expect("should not be null"); - sniff_fmt_datetime(val)? + } + infer::infer_from_values(string_ca, |val| { + datetime_pattern(val, NaiveDateTime::parse_from_str) + .or_else(|| datetime_pattern(val, NaiveDate::parse_from_str)) + }) + .ok_or_else(|| polars_err!(parse_fmt_idk = "datetime"))? }, }; diff --git a/py-polars/tests/unit/operations/namespaces/temporal/test_to_datetime.py b/py-polars/tests/unit/operations/namespaces/temporal/test_to_datetime.py index be59743c5fb1..9c54150a308b 100644 --- a/py-polars/tests/unit/operations/namespaces/temporal/test_to_datetime.py +++ b/py-polars/tests/unit/operations/namespaces/temporal/test_to_datetime.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: from hypothesis.strategies import DrawFn - from polars._typing import TimeUnit + from polars._typing import EngineType, TimeUnit DATE_FORMATS = ["%Y{}%m{}%d", "%d{}%m{}%Y"] @@ -329,3 +329,49 @@ def test_to_datetime_inexact_unicode_multibyte() -> None: assert_frame_equal( out, pl.DataFrame({"a": [datetime(2020, 2, 3, 12, 53, 11), None]}) ) + + +@pytest.mark.parametrize("engine", ["in-memory", "streaming"]) +@pytest.mark.parametrize( + ("method", "dtype"), + [ + ("to_date", pl.Date), + ("to_datetime", pl.Datetime("us")), + ("to_time", pl.Time), + ], +) +def test_strptime_uninferrable_format( + engine: EngineType, method: str, dtype: pl.DataType +) -> None: + lf = pl.LazyFrame({"s": ["not a temporal value"]}) + + non_strict = getattr(pl.col("s").str, method)(strict=False) + assert_frame_equal( + lf.select(non_strict).collect(engine=engine), + pl.DataFrame({"s": [None]}, schema={"s": dtype}), + ) + + strict = getattr(pl.col("s").str, method)(strict=True) + with pytest.raises(ComputeError, match="could not find an appropriate format"): + lf.select(strict).collect(engine=engine) + + +STRPTIME_INFERENCE_CASES = [ + ("to_date", pl.Date, "2020-01-01"), + ("to_datetime", pl.Datetime("us"), "2020-01-01 10:00:00"), + ("to_time", pl.Time, "10:00:00"), +] + + +@pytest.mark.slow +@pytest.mark.parametrize("engine", ["in-memory", "streaming"]) +@pytest.mark.parametrize(("method", "dtype", "good"), STRPTIME_INFERENCE_CASES) +def test_strptime_strict_reports_original_column_name( + engine: EngineType, method: str, dtype: pl.DataType, good: str +) -> None: + n = 200_000 + values = ["not a temporal value"] + [good] * n + + strict = getattr(pl.col("s").str, method)(strict=True) + with pytest.raises(InvalidOperationError, match="in column 's'"): + pl.LazyFrame({"s": values}).select(strict).collect(engine=engine) diff --git a/py-polars/tests/unit/sql/test_cast.py b/py-polars/tests/unit/sql/test_cast.py index 73adb3b6cca3..816704bcf423 100644 --- a/py-polars/tests/unit/sql/test_cast.py +++ b/py-polars/tests/unit/sql/test_cast.py @@ -1,12 +1,17 @@ from __future__ import annotations +from datetime import date, datetime, time from typing import Any import pytest import polars as pl import polars.selectors as cs -from polars.exceptions import InvalidOperationError, SQLInterfaceError +from polars.exceptions import ( + ComputeError, + InvalidOperationError, + SQLInterfaceError, +) from polars.testing import assert_frame_equal @@ -164,20 +169,47 @@ def test_cast() -> None: @pytest.mark.parametrize( - ("values", "cast_op", "error"), + ("values", "cast_op", "exc", "error"), [ - ([1.0, -1.0], "values::uint8", "conversion from `f64` to `u64` failed"), - ([10, 0, -1], "values::uint4", "conversion from `i64` to `u32` failed"), - ([int(1e8)], "values::int1", "conversion from `i64` to `i8` failed"), - (["a", "b"], "values::time", "conversion from `str` to `time` failed"), - (["a", "b"], "values::int4", "conversion from `str` to `i32` failed"), + ( + [1.0, -1.0], + "values::uint8", + InvalidOperationError, + "conversion from `f64` to `u64` failed", + ), + ( + [10, 0, -1], + "values::uint4", + InvalidOperationError, + "conversion from `i64` to `u32` failed", + ), + ( + [int(1e8)], + "values::int1", + InvalidOperationError, + "conversion from `i64` to `i8` failed", + ), + ( + ["a", "b"], + "values::time", + ComputeError, + "could not find an appropriate format to parse times", + ), + ( + ["a", "b"], + "values::int4", + InvalidOperationError, + "conversion from `str` to `i32` failed", + ), ], ) -def test_cast_errors(values: Any, cast_op: str, error: str) -> None: +def test_cast_errors( + values: Any, cast_op: str, exc: type[Exception], error: str +) -> None: df = pl.DataFrame({"values": values}) # invalid CAST should raise an error... - with pytest.raises(InvalidOperationError, match=error): + with pytest.raises(exc, match=error): df.sql(f"SELECT {cast_op} FROM self") # ... or return `null` values if using TRY_CAST @@ -186,6 +218,109 @@ def test_cast_errors(values: Any, cast_op: str, error: str) -> None: assert None in res.to_series() +@pytest.mark.parametrize( + ("sql_type", "dtype", "value", "expected"), + [ + ("date", pl.Date, "2000-02-01", date(2000, 2, 1)), + ( + "timestamp", + pl.Datetime("us"), + "2000-02-01 12:30:00", + datetime(2000, 2, 1, 12, 30), + ), + ( + "datetime", + pl.Datetime("us"), + "2000-02-01 12:30:00", + datetime(2000, 2, 1, 12, 30), + ), + ("time", pl.Time, "12:30:00", time(12, 30)), + ], +) +def test_cast_string_to_temporal( + sql_type: str, dtype: pl.DataType, value: str, expected: Any +) -> None: + df = pl.DataFrame({"s": [value, None]}) + + # a column operand keeps the frame's null; a literal broadcasts to both rows + for operand, rows in (("s", [expected, None]), (f"'{value}'", [expected] * 2)): + for cast_op in ( + f"CAST({operand} AS {sql_type})", + f"TRY_CAST({operand} AS {sql_type})", + f"{operand}::{sql_type}", + ): + res = df.sql(f"SELECT {cast_op} AS x FROM self") + assert_frame_equal(res, pl.DataFrame({"x": rows}, schema={"x": dtype})) + + +@pytest.mark.parametrize( + ("sql_type", "dtype"), + [ + ("date", pl.Date), + ("timestamp", pl.Datetime("us")), + ("time", pl.Time), + ], +) +def test_try_cast_string_to_temporal_nulls(sql_type: str, dtype: pl.DataType) -> None: + df = pl.DataFrame({"s": ["not a temporal value"]}) + + for operand in ("s", "'not a temporal value'"): + res = df.sql(f"SELECT TRY_CAST({operand} AS {sql_type}) AS x FROM self") + assert_frame_equal(res, pl.DataFrame({"x": [None]}, schema={"x": dtype})) + + with pytest.raises(ComputeError, match="could not find an appropriate format"): + df.sql(f"SELECT CAST({operand} AS {sql_type}) AS x FROM self") + + +def test_cast_temporal_to_temporal_is_not_parsed() -> None: + df = pl.DataFrame( + {"dtm": [datetime(2000, 2, 1, 12, 30)]}, + schema={"dtm": pl.Datetime("us")}, + ) + res = df.sql( + """ + SELECT + CAST(dtm AS date) AS d, + CAST(dtm AS time) AS t + FROM self + """ + ) + assert_frame_equal( + res, + pl.DataFrame({"d": [date(2000, 2, 1)], "t": [time(12, 30)]}), + ) + + +def test_cast_string_to_date_in_between() -> None: + df = pl.DataFrame( + {"d": [date(1999, 1, 1), date(1999, 3, 1), date(2000, 1, 1)]}, + ) + res = df.sql( + "SELECT * FROM self WHERE d BETWEEN CAST('1999-02-22' AS date) AND CAST('1999-03-24' AS date)" + ) + assert_frame_equal(res, pl.DataFrame({"d": [date(1999, 3, 1)]})) + + +def test_temporal_in_string_list() -> None: + df = pl.DataFrame( + {"d": [date(1999, 1, 1), date(1999, 3, 1), date(2000, 1, 1)]}, + ) + res = df.sql("SELECT * FROM self WHERE d IN ('1999-03-01', '2000-01-01')") + assert_frame_equal(res, pl.DataFrame({"d": [date(1999, 3, 1), date(2000, 1, 1)]})) + + +def test_try_cast_string_to_temporal_partial() -> None: + df = pl.DataFrame({"s": ["2000-02-01", "nope"]}) + + res = df.sql("SELECT TRY_CAST(s AS date) AS x FROM self") + assert_frame_equal( + res, pl.DataFrame({"x": [date(2000, 2, 1), None]}, schema={"x": pl.Date}) + ) + + with pytest.raises(InvalidOperationError, match=r"conversion .* failed"): + df.sql("SELECT CAST(s AS date) AS x FROM self") + + @pytest.mark.may_fail_cloud # reason: eager construct to_struct @pytest.mark.xfail # this is a construct we cannot deal with anymore def test_cast_json() -> None: